Pytorch实战 |Y4 common.py文件解读

前沿

文件位置：./models/commonpy
该文件是实现YOLO算法中各个模块的地方，如果我们需要修改某一模块（例如C3），那么就需要修改这个文件中对应模块的的定义。由于YOLOv5版本问题，同一个模块可能会看到不同的版本，这都是正常的，以官网为主即可。

• 本周任务：将yolov5s网络模型中的C3模块按照下图方式修改，并跑通yolov5

0、导入需要的包和基本配置

import ast
import contextlib
import json
import math                # 数学函数模块
import platform
import warnings
import zipfile
from collections import OrderedDict, namedtuple
from copy import copy      # 数据拷贝模块，分浅拷贝和深拷贝
from pathlib import Path   # Path将str转换为Path对象，使字符串路径易于操作的模块
from urllib.parse import urlparse

import cv2
import numpy as np          # numpy数组操作模块
import pandas as pd         # pandas数组操作模块
import requests             # Python的HTTP客户端库
import torch                # pytorch深度学习框架
import torch.nn as nn       # 专门为神经网络设计的模块化接口
from IPython.display import display
from PIL import Image       # 图像基础操作模块
from torch.cuda import amp  # 混合精度训练模块

from utils import TryExcept
from utils.dataloaders import exif_transpose, letterbox
from utils.general import (LOGGER, ROOT, Profile, check_requirements, check_suffix, check_version, colorstr,
                           increment_path, is_notebook, make_divisible, non_max_suppression, scale_boxes, xywh2xyxy,
                           xyxy2xywh, yaml_load)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import copy_attr, smart_inference_mode

1、基本组件

1.1 autopad

这个模块可以根据输入的卷积核计算卷积模块所需的pad值。将会用于下面会讲到的 Conv 函数和 Classify 函数中。

def autopad(k, p=None, d=1):  # kernel, padding, dilation
    # Pad to 'same' shape outputs
    if d > 1:
        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p

1.2 Conv

这个函数是整个网络中最基础的组件，由 卷积层 + BN层 + 激活函数组成，具体结构如下

class Conv(nn.Module):
    # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        ''' 在Focus、Bottleneck、BottleneckCSP、C3、SPP、DWConv、TransformerBlock等模块中调用
        Standard convolution : conv + BN + act
        :params c1: 输入的channel值
        :params c2: 输出的channel值
        :params k: 卷积的kernel_size
        :params s: 卷积的stride
        :params p: 卷积的padding，默认是None，可以通过autopad自行计算需要的padding值
        :params g: 卷积的groups数，1就是普通的卷积，>1就是深度可分离卷积
        :params act: 激活函数类型，True就是SiLU()/Swish，False就是不使用激活函数，类型是nn.Module就使用传进来的激活函数类型
        '''
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act 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 forward_fuse(self, x):
        ''' 用于Model类的fuse函数
        融合 Conv + BN 加速推理，一般用于测试/验证阶段
        '''
        return self.act(self.conv(x))

这个类中还有一个特殊函数 forward_fuse ，这是一个前向加速推理模块，在前向传播过程中，通过融合 Conv + BN 层，达到加速推理的作用，一般用于测试或验证阶段

1.3 Focus

Focus模块是作者自己设计出来，为了减少浮点数和提高速度，而不是增加featuremap的，本质就是将图像进行切片，类似于下采样取值，将原图像的宽高信息切分，聚合到channel通道中。结构如下所示：

class Focus(nn.Module):
    # Focus wh information into c-space 把宽度w和高度h的信息整合到c空间中
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        ''' 在yolo.py的parse_model函数中被调用
        理论：从高分辨率图像中，周期性的抽出像素点重构到低分辨率图像中，即将图像相邻的四个位置进行堆叠，
        聚集wh维度信息到c通道中，提高每个点的感受野，并减少原始信息的丢失，该模块的设计主要是减少计算量加快速度。
        先做4个slice，再concat，最后在做Conv
        slice后 (b1,c1,w,h) -> 分成4个slice，每个slice(b,c1,w/2,h/2)
        concat(dim=1)后 4个slice(b,c1,w/2,h/2) -> (b,4c1,w/2,h/2)
        conv后 (b,4c1,w/2,h/2) -> (b,c2,w/2,h/2)
        :params c1: slice后的channel
        :params c2: Focus最终输出的channel
        :params k: 最后卷积的kernel
        :params s: 最后卷积的stride
        :params p: 最后卷积的padding
        :params g: 最后卷积的分组情况，=1普通卷积，>1深度可分离卷积
        :params act: bool激活函数类型，默认True[SiLU()/Swish]，False[不用激活函数]
        '''        
        super().__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act=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))
        # return self.conv(self.contract(x))

1.4 Bottleneck

class Bottleneck(nn.Module):
    # Standard bottleneck  Conv + Conv + shortcut
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        ''' 在BottleneckCSP和yolo.py的parse_model函数中被调用
        :params c1: 第一个卷积的输入channel
        :params c2: 第二个卷积的输入channel
        :params shortcut: bool值，是否有shortcut连接，默认True
        :params g: 卷积分组的个数，=1普通卷积，>1深度可分离卷积
        :params e: expansion ratio，e*c