channel_shuffle代码实现

最新推荐文章于 2025-03-27 20:32:27 发布

原创最新推荐文章于 2025-03-27 20:32:27 发布

· 782 阅读

3 ·

版权

文章标签：

#前端 #python #ide #算法

深度学习专栏收录该内容

40 篇文章

订阅专栏

本文详细介绍了如何在Pytorch中实现ShuffleNetv2模型中的关键层ChannelShuffle，包括图像通道拆分、GConv1处理和Feature合并，以及最终的reshape和transpose操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

结构图，先将输入的图像进行通道拆分为组GConv1，每个GConv1再拆分Feature，每个GConv1的Feature进行合并GConv2，输出Output

输入图像x，拆分为groups个组，每隔组的通道数为channels_per_group

batch_size, num_channels, height, width = x.size()
channels_per_group = num_channels // groups

进行变换

# reshape
# [batch_size, num_channels, height, width] -> [batch_size, groups, channels_per_group, height, width]
x = x.view(batch_size, groups, channels_per_group, height, width)

再将1和2的维度进行调换，就实现了Feature到GConv2

x = torch.transpose(x, 1, 2).contiguous()

全部代码

def channel_shuffle(x: Tensor, groups: int) -> Tensor:

    batch_size, num_channels, height, width = x.size()
    channels_per_group = num_channels // groups

    # reshape
    # [batch_size, num_channels, height, width] -> [batch_size, groups, channels_per_group, height, width]
    x = x.view(batch_size, groups, channels_per_group, height, width)

    x = torch.transpose(x, 1, 2).contiguous()

    # flatten
    x = x.view(batch_size, -1, height, width)

    return x

参考视频：

8.2 使用Pytorch搭建ShuffleNetv2_哔哩哔哩_bilibili