StemBlock pytorch实现

本文mark一下peleeent和yolov5中的实现：

import time
from collections import OrderedDict
import torch.nn.functional as m_func
import torch
import cv2
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class conv_bn_relu(nn.Module):
    def __init__(self, nin, nout, kernel_size, stride, padding, bias=False):
        super(conv_bn_relu, self).__init__()
        self.conv = nn.Conv2d(nin, nout, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
        self.batch_norm = nn.BatchNorm2d(nout)
        self.relu = nn.ReLU6(True)
        # self.relu = mish.Mish()

    def forward(self, x):
        out = self.conv(x)
        out = self.batch_norm(out)
        out = self.rel