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)