mindspore.nn.ChannelShuffle

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

Divide the channels of Tensor whose shape is \((*, C, H, W)\) into \(g\) groups to obtain a Tensor with shape \((*, C \frac g, g, H, W)\), and transpose along the corresponding axis of \(C\), \(\frac{g}{}\) and \(g\) to restore Tensor to the original shape.

Parameters: groups (int) – Number of groups to divide channels in, must be greater than 0. 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

>>> import mindspore as ms
>>> import numpy as np
>>> channel_shuffle = ms.nn.ChannelShuffle(2)
>>> x = ms.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]]]]