mindspore.nn.ChannelShuffle

class mindspore.nn.ChannelShuffle(groups)[source]

Divide the channels in a tensor of shape \((*, C , H, W)\) into g groups and rearrange them as \((*, C \frac g, g, H, W)\), while keeping the original tensor shape.

Parameters: groups (int) – Number of groups to divide channels in. Refer to \(g\).

Inputs:

x (Tensor) - Tensor of shape \((*, C_{in}, H_{in}, W_{in})\).

Outputs:

Tensor, with the same type and shape as the x.

Raises

TypeError – If groups is not an int.
ValueError – If groups is less than 1.
ValueError – If dims of x is less than 3.
ValueError – If number of channels can not be divisible by groups.

Supported Platforms:: Ascend GPU CPU

Examples

>>> channel_shuffle = nn.ChannelShuffle(2)
>>> x = Tensor(np.arange(16).astype(np.int32).reshape(1, 4, 2, 2))
>>> print(x)
[[[[0 1],
   [2 3]],
  [[4 5],
   [6 7]],
  [[8 9],
   [10 11]],
  [[12 13],
   [14 15]],
 ]]
>>> output = channel_shuffle(x)
>>> print(output)
[[[[0 1],
   [2 3]],
  [[8 9],
   [10 11]],
  [[4 5],
   [6 7]],
  [[12 13],
   [14 15]],
 ]]