yolov5 Stemblock

YOLOv5中的Stemblock是网络的起始部分，用于提取输入图像的低级特征。它由一系列卷积层和池化层组成，以逐渐减小特征图的尺寸和通道数。Stemblock的设计旨在捕捉图像的全局信息，并为后续的特征提取过程提供良好的基础。

以下是YOLOv5中Stemblock的示例代码：

import torch
import torch.nn as nn

class Stemblock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Stemblock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool(x)
        return x

# 使用Stemblock
in_channels = 3
out_channels = 64
input = torch.randn(1, in_channels, 416, 416)
stemblock = Stemblock(in_channels, out_channels)
output = stemblock(input)
print(output.shape)  # 输出：torch.Size([1, 64, 208, 208])

在上述代码中，Stemblock接受一个输入张量，经过一系列的卷积、批归一化、ReLU激活和池化操作后，输出一个特征图。这个特征图将作为后续特征提取的输入。