Channel Shuffle类

ShuffleNet 中引入了 channel shuffle, 用来进行不同分组的特征之间的信息流动, 以提高性能。而Detectron2使用的pytorch版本一般较低，没有channel shuffle这个类，因此编写这个轮子用于通道洗牌。实现了与1.11.0官方库相同的结果。

官方文档：ChannelShuffle — PyTorch 1.11.0 documentation

import torch
import torch.nn as nn

class ChannelShuffle(nn.Module):
    def __init__(self, groups):
        super(ChannelShuffle, self).__init__()
        self.groups = groups

    def forward(self, x):
        B, C, H, W = x.shape

        chnls_per_group = C // self.groups
        
        assert C % self.groups == 0

        x = x.view(B, self.groups, chnls_per_group, H, W)  # 通道分组 (B,C,H,W)->(B,group,C,H,W)

        x = torch.transpose(x, 1, 2).contiguous()  # 通道洗牌
        x = x.view(B, -1, H, W)  # 重新展开为(B,C,H,W)
        return x

channel_shuffle = ChannelShuffle(2)

input = torch.randn(1, 4, 2, 2)
print(input)

output = channel_shuffle(input)

print(output)