YOLOv5的输出端head头代码解读

源码分析（Yolo.py中的class Detect）

1、逐份逐份分析版

我们就按他的分法按3个板块来解说。

    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layer
        #yolov5中的anchors（3个，对应Neck出来的那3个输出），初始anchor是由w，h宽高组成，用的是原图的像素尺寸，设置为每层3个，所以共有3 * 3 = 9个
        super().__init__()
        self.nc = nc  # 预测的类的数量
        self.no = nc + 5  
        # 每一个预测框（anchor）输出的数量，对应每种类的置信度（nc），预测框的高宽，中心点坐标，预测框内是否有物体的置信度，共5种信息。
        self.nl = len(anchors)  # 预测层的数量
        self.na = len(anchors[0]) // 2  #预测框的数量

 

# 每一个预测框（anchor）输出的数量，对应每种类的置信度（nc），预测框的高宽，中心点坐标，预测框内是否有物体的置信度，共5种信息。

self.nl = len(anchors) # 预测层的数量

self.na = len(anchors[0]) // 2 #预测框的数量

 

🎈yolov5中的anchors（3个，对应Neck出来的那3个输出），初始anchor是由w，h宽高组成，用的是原图的像素尺寸，设置为每层3个，所以共有3 * 3 = 9个

🎈nc：待预测的类的数量

🎈no：每一个预测框（anchor）输出的数量，对应每种类的置信度（nc），预测框的高宽，中心点坐标，预测框内是否有物体的置信度，共5 + 80 = 85种信息。

🎈nl：预测层的数量

🎈na：预测框的数量

        self.grid = [torch.zeros(1)] * self.nl  # 初始网格，对于每个预测层都有初始网格的生成
        self.anchor_grid = [torch.zeros(1)] * self.nl  # 初始预测框网格
        self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2))  
        # shape(nl,na,2)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  
        # output conv，输出结果：每个预测框的输出结果 * 预测框个数
        self.inplace = inplace  # use in-place ops (e.g. slice assignment)

 

🎈grid：初始网格，对于每个预测层都有初始网格的生成

🎈anchor_grid：初始预测框网格

🎈self.m：output conv，输出结果：每个预测框的输出结果 * 预测框个数

🌳def  _make_grid

准备网格，所有的预测的单位长度都是基于grid层面的而不是原图，并且每一层的grid的尺寸都是不一样的，和每一层输出的尺寸w,h是一样的。

     def _make_grid(self, nx=20, ny=20, i=0):
        #准备网格，所有的预测的单位长度都是基于grid层面的而不是原图，并且每一层的grid的尺寸都是不一样的，和每一层输出的尺寸w,h是一样的。
        d = self.anchors[i].device
        if check_version(torch.__version__, '1.10.0'):  # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
            yv, xv = torch.meshgrid([torch.arange(ny, device=d), torch.arange(nx, device=d)], indexing='ij')
            #torch.meshgrid()生成网格，可以用于生成坐标，尺寸nx * ny；ny范围是竖向坐标；nx范围是横向坐标
        else:
            yv, xv = torch.meshgrid([torch.arange(ny, device=d), torch.arange(nx, device=d)])
        grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()
        anchor_grid = (self.anchors[i].clone() * self.stride[i]) \
            .view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()
        return grid, anchor_grid #制成网格返回

 

🎈基本大框架是先检查版本，针对不同版本进行不同的同款操作

🎈torch.meshgrid()生成网格，可以用于生成坐标，尺寸nx * ny；ny范围是竖向坐标；nx范围是横向坐标

🌳def  forward

主结构是循环每层每层的处理

      def forward(self, x):
        z = []  # inference output
        for i in range(self.nl): #每层循环着处理

 

循环里面进行核心操作

x[i] = self.m[i](x[i])  # conv卷积

 

卷积处理

bs, _, ny, nx = x[i].shape  #bs-batch_size

 

            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            #view()变换形状，数据不变, x(bs,255,20,20) to x(bs,3,85,20,20),将一个预测层里的3个anchor的信息分出来，每个预测框预测信息数量为self.no(这里为85)
            #permute(0, 1, 3, 4, 2)，x[i]有5个维度，（2，3，4）变成（3，4，2）,x(bs,3,85,20,20)to x(bs,3,20,20,85)
            #contiguous()进行一个拷贝

 

🎈view()变换形状，数据不变, x(bs,255,20,20) to x(bs,3,85,20,20),将一个预测层里的3个anchor的信息分出来，每个预测框预测信息数量为self.no(这里为85)

🎈permute(0, 1, 3, 4, 2)，x[i]有5个维度，（2，3，4）变成（3，4，2）,x(bs,3,85,20,20)to x(bs,3,20,20,85)

🎈contiguous()进行一个拷贝

然后又进入一个是否进行训练的框架中（if not self.training:）

            if not self.training:  # inference
                if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
                    #制作第几预测层的网格

 

判断是否需要制作第i预测层的网格

 y = x[i].sigmoid() 
#激活函数，完成逻辑回归的软判决，变量映射到0，1之间的S型函数，所以最后的y就是相对于网格占了几分之几的意思（对center的x,y,w,h都做了归一化处理）

 

激活函数，完成逻辑回归的软判决，变量映射到0，1之间的S型函数，所以最后的y就是相对于网格占了几分之几的意思（对center的x,y,w,h都做了归一化处理）

                if self.inplace:
                    y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    #box center的x,y的预测被乘以2并减去了0.5，让他的预测范围变成（-0.5，1.5）就是能跨半个网格预测
                    #然后加上self.grid[i]，就是加上网格的宽度/高度
                    #最后乘上self.stride[i]，就是步长，定位到原先预测的那个点
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                    #对预测框高宽的处理
                else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                    xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, y[..., 4:]), -1)

 

这里分析self.inplace的情况：

🎈第一句是对中心点x，y的操作。

（1）box center的x,y的预测被乘以2并减去了0.5，让他的预测范围变成（-0.5，1.5）就是能跨半个网格预测

（2）然后加上self.grid[i]，就是加上网格的宽度/高度

（3）最后乘上self.stride[i]，就是步长，定位到原先预测的那个点

🎈第二句是对预测框的宽高进行操作。

z.append(y.view(bs, -1, self.no))

 

#将结果填入z:(第几层的预测层),(预测出的center的x，y，以及预测框的w和h),(对应的85种信息--这里个人认为这85种信息中的x,y,w,h不会再用到，主要是取出置信度信息)

将结果填入z:(第几层的预测层),(预测出的center的x，y，以及预测框的w和h),(对应的85种信息--这里个人认为这85种信息中的x,y,w,h不会再用到，主要是取出置信度信息)

return x if self.training else (torch.cat(z, 1), x)

 

最后来个判断后return，如果还要训练就是返回x，如果训练完毕，那就返回(torch.cat(z, 1), x)

2、代码注释分析一体化

class Detect(nn.Module):
    stride = None  # strides computed during build
    onnx_dynamic = False  # ONNX export parameter
 
    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):  # detection layer
        #yolov5中的anchors（3个，对应Neck出来的那3个输出），初始anchor是由w，h宽高组成，用的是原图的像素尺寸，设置为每层3个，所以共有3 * 3 = 9个
        super().__init__()
        self.nc = nc  # 预测的类的数量
        self.no = nc + 5  
        # 每一个预测框（anchor）输出的数量，对应每种类的置信度（nc），预测框的高宽，中心点坐标，预测框内是否有物体的置信度，共5种信息。
        self.nl = len(anchors)  # 预测层的数量
        self.na = len(anchors[0]) // 2  #预测框的数量
        self.grid = [torch.zeros(1)] * self.nl  # 初始网格，对于每个预测层都有初始网格的生成
        self.anchor_grid = [torch.zeros(1)] * self.nl  # 初始预测框网格
        self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2))  
        # shape(nl,na,2)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  
        # output conv，输出结果：每个预测框的输出结果 * 预测框个数
        self.inplace = inplace  # use in-place ops (e.g. slice assignment)
 
    def forward(self, x):
        z = []  # inference output
        for i in range(self.nl): #每层循环着处理
            x[i] = self.m[i](x[i])  # conv卷积
            bs, _, ny, nx = x[i].shape  
            #bs第几个预测层的意思吧
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            #view()变换形状，数据不变, x(bs,255,20,20) to x(bs,3,85,20,20),将一个预测层里的3个anchor的信息分出来，每个预测框预测信息数量为self.no(这里为85)
            #permute(0, 1, 3, 4, 2)，x[i]有5个维度，（2，3，4）变成（3，4，2）,x(bs,3,85,20,20)to x(bs,3,20,20,85)
            #contiguous()进行一个拷贝
            if not self.training:  # inference
                if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
                    #制作第几预测层的网格
 
                y = x[i].sigmoid()
                #激活函数，完成逻辑回归的软判决，变量映射到0，1之间的S型函数，所以最后的y就是相对于网格占了几分之几的意思（对center的x,y,w,h都做了归一化处理）
                if self.inplace:
                    y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    #box center的x,y的预测被乘以2并减去了0.5，让他的预测范围变成（-0.5，1.5）就是能跨半个网格预测
                    #然后加上self.grid[i]，就是加上网格的宽度/高度
                    #最后乘上self.stride[i]，就是步长，定位到原先预测的那个点
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                    #对预测框高宽的处理
                else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                    xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, y[..., 4:]), -1)
                z.append(y.view(bs, -1, self.no))
                #将结果填入z:(第几层的预测层),(预测出的center的x，y，以及预测框的w和h),(对应的85种信息--这里个人认为这85种信息中的x,y,w,h不会再用到，主要是取出置信度信息)
 
        return x if self.training else (torch.cat(z, 1), x)
 
    def _make_grid(self, nx=20, ny=20, i=0):
        #准备网格，所有的预测的单位长度都是基于grid层面的而不是原图，并且每一层的grid的尺寸都是不一样的，和每一层输出的尺寸w,h是一样的。
        d = self.anchors[i].device
        if check_version(torch.__version__, '1.10.0'):  # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
            yv, xv = torch.meshgrid([torch.arange(ny, device=d), torch.arange(nx, device=d)], indexing='ij')
            #torch.meshgrid()生成网格，可以用于生成坐标，尺寸nx * ny；ny范围是竖向坐标；nx范围是横向坐标
        else:
            yv, xv = torch.meshgrid([torch.arange(ny, device=d), torch.arange(nx, device=d)])
        grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()
        anchor_grid = (self.anchors[i].clone() * self.stride[i]) \
            .view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()
        return grid, anchor_grid #制成网格返回

 

---

欢迎大家在评论区批评指正，谢谢大家~