Shortcuts

ChannelShuffle

class torch.nn.ChannelShuffle(groups)[source]

Divides and rearranges the channels in a tensor.

This operation divides the channels in a tensor of shape (N,C,)(N, C, *) into g groups as (N,Cg,g,)(N, \frac{C}{g}, g, *) and shuffles them, while retaining the original tensor shape in the final output.

Parameters

groups (int) – number of groups to divide channels in.

Examples:

>>> channel_shuffle = nn.ChannelShuffle(2)
>>> input = torch.arange(1, 17, dtype=torch.float32).view(1, 4, 2, 2)
>>> input
tensor([[[[ 1.,  2.],
          [ 3.,  4.]],
         [[ 5.,  6.],
          [ 7.,  8.]],
         [[ 9., 10.],
          [11., 12.]],
         [[13., 14.],
          [15., 16.]]]])
>>> output = channel_shuffle(input)
>>> output
tensor([[[[ 1.,  2.],
          [ 3.,  4.]],
         [[ 9., 10.],
          [11., 12.]],
         [[ 5.,  6.],
          [ 7.,  8.]],
         [[13., 14.],
          [15., 16.]]]])

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources