yolov5focus代码

以下是 YOLOv5 Focus 的代码实现：

import torch.nn as nn

class Focus(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1):
        super(Focus, self).__init__()
        self.conv = nn.Conv2d(in_channels * 4, out_channels, kernel_size, stride=1, padding=kernel_size // 2)

    def forward(self, x):
        # x shape: [batch_size, in_channels, height, width]
        # split tensor along channel axis
        x = torch.cat([x[:, :, ::2, ::2], x[:, :, 1::2, ::2], x[:, :, ::2, 1::2], x[:, :, 1::2, 1::2]], dim=1)
        # apply convolution
        x = self.conv(x)
        return x

其中，`in_channels` 表示输入张量的通道数，`out_channels` 表示输出张量的通道数，`kernel_size` 表示卷积核大小。

该模块的作用是将输入张量进行通道拆分，然后在通道维度上进行卷积操作，最后将结果进行拼接并输出。