class DepthwiseSeparableConv(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
super(DepthwiseSeparableConv, self).__init__()
# 深度卷积层
self.depthwise = nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size,
stride, padding, groups=in_channels),
nn.BatchNorm2d(in_channels),
# activation_layer
nn.LeakyReLU(0.1, inplace=True)
)
# 逐点卷积层
self.pointwise = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1),
nn.BatchNorm2d(out_channels),
# activation_layer
nn.LeakyReLU(0.1, inplace=True)
)
def forward(self, x):
x = self.depthwise(x)
x = self.pointwise(x)
return x
深度可分离卷积代码
最新推荐文章于 2024-04-30 03:13:08 发布