| 8 | |
| 9 | |
| 10 | class ShuffleBlock(nn.Module): |
| 11 | def __init__(self, groups): |
| 12 | super(ShuffleBlock, self).__init__() |
| 13 | self.groups = groups |
| 14 | |
| 15 | def forward(self, x): |
| 16 | '''Channel shuffle: [N,C,H,W] -> [N,g,C/g,H,W] -> [N,C/g,g,H,w] -> [N,C,H,W]''' |
| 17 | N,C,H,W = x.size() |
| 18 | g = self.groups |
| 19 | return x.view(N,g, int(C/g),H,W).permute(0,2,1,3,4).contiguous().view(N,C,H,W) |
| 20 | |
| 21 | |
| 22 | class Bottleneck(nn.Module): |