ShuffleNet v2网络结构复现
from torch import nn from torch.nn import functional import torch from torchsummary import summary # ---------------------------- ShuffleBlock start ------------------------------- # 通道重排，跨group信息交流 def channel_shuffle(x, groups): batchsize, num_channels, height, width = x.data.size() channels_per_group = num_channels // groups # reshape x = x.view(batchsize, groups, channels_per_group, height, width) x = torch.transpose(x, 1, 2).contiguous() # flatten x = x.view(batchsize, -1, height, width) return x class CBRM(nn.Module): def __init__(self, c1, c2): # ch_in, ch_out super(CBRM, self).__init__() self.conv = nn.Sequential( nn.Conv2d(c1, c2, kernel_size=3, stride=2, padding=1, bias=False), nn.BatchNorm2d(c2), nn.ReLU(inplace=True), ) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) def forward(self, x): return self.maxpool(self.conv(x)) class Shuffle_Block(nn.Module): def __init__(self, inp, oup, stride): super(Shuffle_Block, self).__init__() if not (1 <= stride <= 3): raise ValueError('illegal stride value') self.stride = stride branch_features = oup // 2 assert (self.stride != 1) or (inp == branch_features << 1) if self.stride > 1: self.branch1 = nn.Sequential( self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1), nn.BatchNorm2d(inp), nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), nn.ReLU(inplace=True), ) self.branch2 = nn.Sequential( nn.Conv2d(inp if (self.stride > 1) else branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), nn.ReLU(inplace=True), self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1), nn.BatchNorm2d(branch_features), nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), nn.ReLU(inplace=True), ) @staticmethod def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False): return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i) def forward(self, x): if self.stride == 1: x1, x2 = x.chunk(2, dim=1) # 按照维度1进行split out = torch.cat((x1, self.branch2(x2)), dim=1) else: out = torch.cat((self.branch1(x), self.branch2(x)), dim=1) out = channel_shuffle(out, 2) return out class ShuffleNetV2(nn.Module): def __init__(self): super(ShuffleNetV2, self).__init__() self.MobileNet_01 = nn.Sequential( CBRM(3, 32), # 160x160 Shuffle_Block(32, 128, 2), # 80x80 Shuffle_Block(128, 128, 1), # 80x80 Shuffle_Block(128, 256, 2), # 40x40 Shuffle_Block(256, 256, 1), # 40x40 Shuffle_Block(256, 512, 2), # 20x20 Shuffle_Block(512, 512, 1), # 20x20 ) def forward(self, x): x = self.MobileNet_01(x) return x if __name__ == '__main__': shufflenetv2 = ShuffleNetV2() device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') inputs = shufflenetv2.to(device) summary(inputs, (3, 640, 640), batch_size=1, device="cuda") # 分别是输入数据的三个维度 #print(shufflenetv2)
---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [1, 32, 320, 320] 864 BatchNorm2d-2 [1, 32, 320, 320] 64 ReLU-3 [1, 32, 320, 320] 0 MaxPool2d-4 [1, 32, 160, 160] 0 CBRM-5 [1, 32, 160, 160] 0 Conv2d-6 [1, 32, 80, 80] 288 BatchNorm2d-7 [1, 32, 80, 80] 64 Conv2d-8 [1, 64, 80, 80] 2,048 BatchNorm2d-9 [1, 64, 80, 80] 128 ReLU-10 [1, 64, 80, 80] 0 Conv2d-11 [1, 64, 160, 160] 2,048 BatchNorm2d-12 [1, 64, 160, 160] 128 ReLU-13 [1, 64, 160, 160] 0 Conv2d-14 [1, 64, 80, 80] 576 BatchNorm2d-15 [1, 64, 80, 80] 128 Conv2d-16 [1, 64, 80, 80] 4,096 BatchNorm2d-17 [1, 64, 80, 80] 128 ReLU-18 [1, 64, 80, 80] 0 Shuffle_Block-19 [1, 128, 80, 80] 0 Conv2d-20 [1, 64, 80, 80] 4,096 BatchNorm2d-21 [1, 64, 80, 80] 128 ReLU-22 [1, 64, 80, 80] 0 Conv2d-23 [1, 64, 80, 80] 576 BatchNorm2d-24 [1, 64, 80, 80] 128 Conv2d-25 [1, 64, 80, 80] 4,096 BatchNorm2d-26 [1, 64, 80, 80] 128 ReLU-27 [1, 64, 80, 80] 0 Shuffle_Block-28 [1, 128, 80, 80] 0 Conv2d-29 [1, 128, 40, 40] 1,152 BatchNorm2d-30 [1, 128, 40, 40] 256 Conv2d-31 [1, 128, 40, 40] 16,384 BatchNorm2d-32 [1, 128, 40, 40] 256 ReLU-33 [1, 128, 40, 40] 0 Conv2d-34 [1, 128, 80, 80] 16,384 BatchNorm2d-35 [1, 128, 80, 80] 256 ReLU-36 [1, 128, 80, 80] 0 Conv2d-37 [1, 128, 40, 40] 1,152 BatchNorm2d-38 [1, 128, 40, 40] 256 Conv2d-39 [1, 128, 40, 40] 16,384 BatchNorm2d-40 [1, 128, 40, 40] 256 ReLU-41 [1, 128, 40, 40] 0 Shuffle_Block-42 [1, 256, 40, 40] 0 Conv2d-43 [1, 128, 40, 40] 16,384 BatchNorm2d-44 [1, 128, 40, 40] 256 ReLU-45 [1, 128, 40, 40] 0 Conv2d-46 [1, 128, 40, 40] 1,152 BatchNorm2d-47 [1, 128, 40, 40] 256 Conv2d-48 [1, 128, 40, 40] 16,384 BatchNorm2d-49 [1, 128, 40, 40] 256 ReLU-50 [1, 128, 40, 40] 0 Shuffle_Block-51 [1, 256, 40, 40] 0 Conv2d-52 [1, 256, 20, 20] 2,304 BatchNorm2d-53 [1, 256, 20, 20] 512 Conv2d-54 [1, 256, 20, 20] 65,536 BatchNorm2d-55 [1, 256, 20, 20] 512 ReLU-56 [1, 256, 20, 20] 0 Conv2d-57 [1, 256, 40, 40] 65,536 BatchNorm2d-58 [1, 256, 40, 40] 512 ReLU-59 [1, 256, 40, 40] 0 Conv2d-60 [1, 256, 20, 20] 2,304 BatchNorm2d-61 [1, 256, 20, 20] 512 Conv2d-62 [1, 256, 20, 20] 65,536 BatchNorm2d-63 [1, 256, 20, 20] 512 ReLU-64 [1, 256, 20, 20] 0 Shuffle_Block-65 [1, 512, 20, 20] 0 Conv2d-66 [1, 256, 20, 20] 65,536 BatchNorm2d-67 [1, 256, 20, 20] 512 ReLU-68 [1, 256, 20, 20] 0 Conv2d-69 [1, 256, 20, 20] 2,304 BatchNorm2d-70 [1, 256, 20, 20] 512 Conv2d-71 [1, 256, 20, 20] 65,536 BatchNorm2d-72 [1, 256, 20, 20] 512 ReLU-73 [1, 256, 20, 20] 0 Shuffle_Block-74 [1, 512, 20, 20] 0 ================================================================ Total params: 445,824 Trainable params: 445,824 Non-trainable params: 0 ---------------------------------------------------------------- Input size (MB): 4.69 Forward/backward pass size (MB): 270.31 Params size (MB): 1.70 Estimated Total Size (MB): 276.70 ----------------------------------------------------------------
ShuffleNetV2( (MobileNet_01): Sequential( (0): CBRM( (conv): Sequential( (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) ) (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) ) (1): Shuffle_Block( (branch1): Sequential( (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False) (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (4): ReLU(inplace=True) ) (branch2): Sequential( (0): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=64, bias=False) (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (5): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) (6): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (7): ReLU(inplace=True) ) ) (2): Shuffle_Block( (branch2): Sequential( (0): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=64, bias=False) (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (5): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) (6): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (7): ReLU(inplace=True) ) ) (3): Shuffle_Block( (branch1): Sequential( (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=128, bias=False) (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (4): ReLU(inplace=True) ) (branch2): Sequential( (0): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=128, bias=False) (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (5): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) (6): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (7): ReLU(inplace=True) ) ) (4): Shuffle_Block( (branch2): Sequential( (0): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128, bias=False) (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (5): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) (6): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (7): ReLU(inplace=True) ) ) (5): Shuffle_Block( (branch1): Sequential( (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=256, bias=False) (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) (3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (4): ReLU(inplace=True) ) (branch2): Sequential( (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=256, bias=False) (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (5): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) (6): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (7): ReLU(inplace=True) ) ) (6): Shuffle_Block( (branch2): Sequential( (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU(inplace=True) (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=256, bias=False) (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (5): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) (6): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (7): ReLU(inplace=True) ) ) ) )
