apple公司提出的MobileVit模型，自己根据论文复现了一下模型
import torch.nn as nn
def conv3x3(in_ch: int, out_ch: int,group:int = 1, stride: int = 1):
return nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=3, stride=stride, padding=1,groups=group);
def conv1x1(in_ch: int, out_ch: int,group:int = 1):
return nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=1,groups=group);
class MobileNetV2_Block(nn.Module):
def __init__(self,in_ch: int, out_ch: int,stride: int = 1):
super(MobileNetV2_Block, self).__init__()
self.conv1 = conv1x1(in_ch,in_ch * 6,group=in_ch * 6)
self.conv2 = conv3x3(in_ch * 6,in_ch * 6,stride=stride)
self.conv3 = conv1x1(in_ch * 6,out_ch,group=out_ch)
def forward(self,x):
x = self.conv1(x)
x = nn.ReLU6(x)
x = self.conv2(x)
x = nn.ReLU6(x)
x = self.conv3(x)
return x
class MobileVit_Block(nn.Module):
def __init__(self,in_ch: int, out_ch: int,d_model:int,nhead:int = 2,
num_encoder_layers:int = 6,num_decoder_layers:int = 6,
dim_feedforward:int = 2048):
super(MobileVit_Block, self).__init__()
self.d_model = d_model
self.conv1 = conv3x3(in_ch,out_ch,group=out_ch)
self.conv2 = conv1x1(out_ch,d_model)
self.transformer = nn.Transformer(d_model=d_model,nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward)
self.conv3 = conv1x1(d_model,out_ch)
self.conv4 = conv1x1(out_ch, out_ch * 3)
self.conv5 = conv3x3(out_ch * 3, out_ch)
def forward(self, x):
h,w = x.shape()[2:]
x = self.conv1(x)
x = self.conv2(x)
x = x.permute(0,2,3,1)
x = x.view(-1,h*w,self.d_model)
x = self.transformer(x)
x = x.view(-1,h,w,self.d_model)
x = x.permute(0, 3, 1, 2)
x = self.conv3(x)
x = self.conv4(x)
x = self.conv5(x)
return x
class Mobile_Vit(nn.Module):
def __init__(self,cnn_block,trans_block,in_ch, out_ch):
super(Mobile_Vit, self).__init__()
self.conv1 = conv3x3(in_ch,in_ch,stride=2)
self.mv1 = self._make_cnn_layer(cnn_block,in_ch,in_ch)
self.mv2 = self._make_cnn_layer(cnn_block,in_ch,in_ch,stride=2)
self.mv3 = self._make_cnn_layer(cnn_block,in_ch,in_ch,2)
self.mv4 = self._make_cnn_layer(cnn_block,in_ch,in_ch,stride=2)
self.mt1 = self._make_trans_layer(trans_block,in_ch,in_ch,512,2)
self.mv5 = self._make_cnn_layer(cnn_block,in_ch,in_ch,stride=2)
self.mt2 = self._make_trans_layer(trans_block,in_ch,in_ch,512,4)
self.mv6 = self._make_cnn_layer(cnn_block, in_ch, in_ch, stride=2)
self.mt3 = self._make_trans_layer(trans_block,in_ch,in_ch,512,3)
self.conv2 = conv1x1(in_ch, out_ch)
self.gap = nn.AdaptiveAvgPool2d(1)
self.flat = nn.Flatten()
def _make_cnn_layer(block, in_ch, out_ch, blocks, stride=1):
layers = []
for i in range(0, blocks):
layers.append(block(in_ch, out_ch, stride))
return nn.Sequential(*layers)
def _make_trans_layer(block, in_ch, out_ch,d_model, blocks):
layers = []
for i in range(0, blocks):
layers.append(block(in_ch, out_ch,d_model))
return nn.Sequential(*layers)
def forward(self,x):
x = self.conv1(x)
x = self.mv1(x)
x = self.mv2(x)
x = self.mv3(x)
x = self.mv4(x)
x = self.mt1(x)
x = self.mv5(x)
x = self.mt2(x)
x = self.mv6(x)
x = self.mt3(x)
x = self.conv2(x)
x = self.flat(self.gap(x))
return x欢迎大家批评指正，论文链接：https://arxiv.org/pdf/2110.02178.pdf
