import torch
import torch.nn as nn
# 深度可分离卷积
class Depth_Separable_Conv(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.stride = stride
self.dw_conv = self.depth_wise_conv() # 逐通道卷积:groups=in_channels
self.pw_conv = self.point_wise_conv() # 逐点卷积,即普通1x1卷积
# 逐通道卷积:groups=in_channels
def depth_wise_conv(self):
return nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.in_channels, kernel_size=3, stride=self.stride,
padding=1, groups=self.in_channels, bias=False),
nn.BatchNorm2d(self.in_channels),
nn.ReLU()
)
# 逐点卷积,即普通1x1卷积
def point_wise_conv(self):
return nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=1, stride=1, padding=0,
groups=1, bias=False),
nn.BatchNorm2d(self.out_channels),
nn.ReLU()
)
def forward(self, x):
x = self.dw_conv(x)
x = self.pw_conv(x)
return x
class MobileNetV1(nn.Module):
def __init__(self, num_classes=1000, img_channels=3):
super().__init__()
# 输出类别数
self.num_classes = num_classes
# 输入图片的通道数
self.img_channels = img_channels
# 网络结构
self.m = nn.Sequential(
nn.Conv2d(in_channels=self.img_channels, out_channels=32, kernel_size=3, stride=2, padding=1, groups=1,
bias=False),
nn.BatchNorm2d(32),
nn.ReLU(),
# Depth_Separable_Conv的参数顺序:in_channels,out_channels,stride
Depth_Separable_Conv(32, 64, 1),
Depth_Separable_Conv(64, 128, 2),
Depth_Separable_Conv(128, 128, 1),
Depth_Separable_Conv(128, 256, 2),
Depth_Separable_Conv(256, 256, 1),
Depth_Separable_Conv(256, 512, 2),
# 重复5次
Depth_Separable_Conv(512, 512, 1),
Depth_Separable_Conv(512, 512, 1),
Depth_Separable_Conv(512, 512, 1),
Depth_Separable_Conv(512, 512, 1),
Depth_Separable_Conv(512, 512, 1),
Depth_Separable_Conv(512, 1024, 2),
Depth_Separable_Conv(1024, 1024, 2),
nn.AdaptiveAvgPool2d((1, 1)),
nn.Flatten(start_dim=1),
nn.Linear(in_features=1024, out_features=self.num_classes)
)
def forward(self, x):
x = self.m(x)
return x
if __name__=='__main__':
model = MobileNetV1()
print(model)
7、mobilenetv1
最新推荐文章于 2023-07-17 11:26:23 发布
关键词由CSDN通过智能技术生成