（beta）torch_npu.contrib.module.ChannelShuffle

接口原型

ChannelShuffle(nn.Module):

功能描述

应用NPU兼容的通道shuffle操作。为避免NPU上效率不高的连续操作，我们用相同语义重写替换原始操作。以下两个不连续操作已被替换：transpose和chunk。

参数说明

• Input (Tensor) - 输入张量。 (N, C_{in}, L_{in}), (N, C_{in}, L_{in})。
• in_channels (Int) - 输入张量中的通道总数。
• groups (Int，默认值为2) - shuffle组数。
• split_shuffle (Bool，默认值为True) - Shuffle后是否执行chunk操作。默认值：True。

输出说明

Output (Tensor) - 输出张量(N, C_{out}, L_{out})。

约束说明

只实现了group=2，请自行修改其他group场景。

调用示例

from torch_npu.contrib.module import ChannelShuffle
m = ChannelShuffle(64, split_shuffle=True)

   >>> x1 = torch.randn(2,32,7,7)
   >>> x2 = torch.randn(2,32,7,7)
   >>> m = ChannelShuffle(64, split_shuffle=True)
   >>> output = m(x1, x2)