| 50 | |
| 51 | |
| 52 | class ShuffleNet(nn.Module): |
| 53 | def __init__(self, cfg, channels=3, num_classes=10, dropout=False): |
| 54 | super(ShuffleNet, self).__init__() |
| 55 | # print (cfg, channels, num_classes) |
| 56 | out_planes = cfg['out_planes'] |
| 57 | num_blocks = cfg['num_blocks'] |
| 58 | groups = cfg['groups'] |
| 59 | |
| 60 | self.conv1 = nn.Conv2d(channels, 24, kernel_size=1, bias=False) |
| 61 | self.bn1 = nn.BatchNorm2d(24) |
| 62 | self.in_planes = 24 |
| 63 | self.layer1 = self._make_layer(out_planes[0], num_blocks[0], groups) |
| 64 | self.layer2 = self._make_layer(out_planes[1], num_blocks[1], groups) |
| 65 | self.layer3 = self._make_layer(out_planes[2], num_blocks[2], groups) |
| 66 | self.linear = nn.Linear(out_planes[2], num_classes) |
| 67 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) |
| 68 | self.dropout = dropout |
| 69 | |
| 70 | def _make_layer(self, out_planes, num_blocks, groups): |
| 71 | layers = [] |
| 72 | for i in range(num_blocks): |
| 73 | stride = 2 if i == 0 else 1 |
| 74 | cat_planes = self.in_planes if i == 0 else 0 |
| 75 | layers.append(Bottleneck(self.in_planes, out_planes-cat_planes, stride=stride, groups=groups)) |
| 76 | self.in_planes = out_planes |
| 77 | return nn.Sequential(*layers) |
| 78 | |
| 79 | def forward(self, x, intermediate=False): |
| 80 | out0 = F.relu(self.bn1(self.conv1(x))) |
| 81 | out1 = self.layer1(out0) |
| 82 | out2 = self.layer2(out1) |
| 83 | out3 = self.layer3(out2) |
| 84 | out = F.avg_pool2d(out3, 4) |
| 85 | out = self.avgpool(out) |
| 86 | e1 = out.view(out.size(0), -1) |
| 87 | if self.dropout: |
| 88 | e1 = F.dropout(e1, p=0.5, training=True) |
| 89 | out = self.linear(e1) |
| 90 | |
| 91 | in_values = [out0, out1, out2, out3] |
| 92 | if intermediate==True: |
| 93 | return out, e1, in_values |
| 94 | else: |
| 95 | return out, e1 |
| 96 | |
| 97 | |
| 98 | def ShuffleNetG2(channels=3,num_classes=10, dropout=False): |
no outgoing calls
no test coverage detected