Notes on ShuffleNet

ShuffleNet 主要是对 feature map 的 channel 做 shuffle 操作. 另外, 还借鉴了近期用的比较多的 depthwise separable convolution 的思想. 另外在 alexnet 中因为当时显存不足提出的 group convolution, 随后随着显存的不断增大逐步被弃用的 group convolution 在 ShuffleNet 中配合 shuffle layer 操作大量使用.

在 group convolution 中, 每一个 feature map 的特征提取是分 group 学习的, 这样导致提取的特征比较 local, 通过 shuffle 操作, 把每个 group 的信息进行了交换.
ShuffleNet 的设计沿用了最近流行的做法, 先设计一个 building block, 通过 building block 的串联组合形成最终的网络.

ShuffleNet unit 改进了 ResNet unit 的设计. 具体来说体现在 3 个方面:

  1. residual branch 上使用 \(3 \times 3\) 的 depthwise convolution 操作取代了 ResNet unit 中的 \(3 \times 3\) convolution
  2. 使用 pointwise group convolution 取代了 ResNet 中的第一个 \(1 times 1\) convolution, 紧接着增加了 shuffle channel 操作.
  3. 可以看到最后还有一个 pointwise group convolution, 作用是完成和 shortcut branch 的维数匹配.

以上设计的输入和输出保持了 feature map 的 shape 完全相同, 最后进行了一个 elementwise 的 add 操作.

在维数变化(w 和 h 减半)的情况下, unit 的设计稍有不同, 如上(c).

  1. shortcut branch 上使用 \(3 \times 3\) 的 average pooling, stride=2, 从而完成了 w 和 h 减半的目的
  2. residual branch 上的 \(3 \times 3\) convolution 的 stride = 2, 完成 w 和 h 减半.
  3. shortcut branch 和 residual branch 的融合使用 concat 操作而不是 add 操作.(在 densenet 中也是使用 concat 而不是 ResNet 中的 add, 我个人也比较喜欢 concat 操作, 但是, concat 操作的缺点是feature map 的 depth 增加比较快, 导致计算量增加.)

ShuffleNet 的实现

ShuffleNet 的 shuffle 并不是完全的随机, 在具体实现时实际上是:

  1. 对 feature map 进行转置
  2. 对转置后的 feature map 进行均匀分组, 分组的数量和转置之前的 feature map 的分组数量相同