yolov5 focus

import torch
import torch.nn as nn
import numpy as np
import pandas as pd
def autopad(k, p=None):  # kernel, padding
    # Pad to 'same'
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p
class Conv(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))
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)
        # self.contract = Contract(gain=2)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
x = np.array([
    [
        [
        [1,2,3,4],
        [2,3,4,5],
        [3,4,5,6],
        [1,2,3,4],
        [2,3,4,5,],
        [3,4,5,6]
        ]
    ]
],dtype = np.float32)
print(x.shape)
x = np.zeros((1,1,4,4), dtype = np.float32)
print(x.shape)
foc = Focus(1,4)
x1 = torch.from_numpy(x)
print(x1)
print(x1[..., ::2, ::2])
print(x1[..., 1::2, ::2])
print(x1[..., ::2, 1::2])
print(torch.cat([x1[..., ::2, ::2], x1[..., 1::2, ::2], x1[..., ::2, 1::2], x1[..., 1::2, 1::2]], 1).shape)
y = foc.forward(x1)
print(y)
print(y.shape)