channel_shuffle

→Tensor

channel_shuffle(x: Tensor, groups: int)

source

Group-then-transpose channel rearrangement (ShuffleNet).

Splits the channel axis into groups equal-sized groups, transposes the group and intra-group axes, and flattens back. This re-orders channels so that, after a subsequent grouped convolution, information from each input group reaches every output group — enabling cheap cross-group information flow that grouped convolutions alone cannot provide.

Conceptually:

(N, g \cdot c, H, W) \xrightarrow{\mathrm{reshape}} (N, g, c, H, W) \xrightarrow{\mathrm{transpose}} (N, c, g, H, W) \xrightarrow{\mathrm{reshape}} (N, g \cdot c, H, W)

where g = groups and c = C / groups.

Parameters

xTensor

Input tensor of shape (N, C, *spatial) (at least 2-D). C must be divisible by groups.

groupsint

Number of channel groups

g

Returns

Tensor

Channel-shuffled tensor of the same shape as x.

Notes

Introduced in Zhang et al., 2018 — "ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices". The operation is a pure reshape + permute, so it adds no FLOPs and very little memory traffic.

Examples

>>> import lucid
>>> from lucid.nn.functional import channel_shuffle
>>> x = lucid.arange(24).reshape(1, 6, 2, 2).astype(lucid.float32)
>>> y = channel_shuffle(x, groups=3)
>>> y.shape
(1, 6, 2, 2)