🚀 Feature

Modify Focus Layer into row major to be compatible with tf.space_to_depth.

Just change the blocks order:
from : torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
to : torch.cat([x[..., ::2, ::2], x[..., ::2, 1::2], x[..., 1::2, ::2], x[..., 1::2, 1::2]], 1)

Motivation

In model/common.py, the Focus Layer is defined in Pytorch as following:

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        # original 
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
        # suggestion 
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., ::2, 1::2], x[..., 1::2, ::2], x[..., 1::2, 1::2]], 1))

And @bonlime posted a brief answer to What's the Focus layer? #207:

check TResNet paper. p2. They call it SpaceToDepth

In the TResNet paper, p2.1 We wanted to create a fast, seamless stem layer, with little information loss as possible, and let the simple well designed residual blocks do all the actual processing work. The stem sole functionality should be to downscale the input resolution to match the rest of the architecture, e.g., by a factor of 4. We met these goals by using a dedicated SpaceToDepth transformation layer [32], that rearranges blocks of spatial data into depth. The SpaceToDepth transformation layer is followed by simple 1x1 convolution to match the number of wanted channels.

That to say, the focus layer is to fast download the input resolution by rearanging blocks of spatial data into depth, and change the feature channels generally by 1x1 conv.

And there is an op SpaceToDepth (tf.space_to_depth, tf2.nn.space_to_depth) to rearranges blocks of spatial data.

The Fcous layer use torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1).

Then we compare:

(0) input

[[[[0 1]
   [2 3]]]]

(1) by Focus torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)

[[[[0]]
  [[2]]
  [[1]]
  [[3]]]]

(2) by tensorflow

[[[[0]]
  [[1]]
  [[2]]
  [[3]]]]

(3) modify Focus torch.cat([x[..., ::2, ::2], x[..., ::2, 1::2], x[..., 1::2, ::2], x[..., 1::2, 1::2]], 1)

[[[[0]]
  [[1]]
  [[2]]
  [[3]]]]

So, just modify the order of the blocks, we can make it compatible tensorflow SpaceToDepth op.
It will make the model be more likely to transport into tensorflow.