rock带你读CornerNet-lite系列源码（二）

前言

接上回rock带你读CornerNet-lite系列源码（一）, 前篇文章介绍了项目代码的总体架构，和训练时的调用关系，数据流传递到了 py_utils.py下的model定义部分，本篇主要介绍（一）py_utils.py下的三个文件，模型定义。（二）sample下的三个文件，构建Groundtruth，encode方式。

CorNerNet 结构

阅读源码最好的方式是按照组件解读，这里强烈建议看下：Hourglass网络的理解和代码分析
好的，这里默认看懂了哈，下面看这段代码就比较好理解了。

import torch
import torch.nn as nn

from .py_utils import TopPool, BottomPool, LeftPool, RightPool
#作者定义的C++4个扩展POOL操作
from .py_utils.utils import convolution, residual, corner_pool
from .py_utils.losses import CornerNet_Loss
from .py_utils.modules import hg_module, hg, hg_net

def make_pool_layer(dim):
    return nn.Sequential()

#重复的残差模块，不会改变特征图的大小，但会改变channel的数量，
#即（B，N， W，H）这个操作只会改变N
def make_hg_layer(inp_dim, out_dim, modules):   
    layers  = [residual(inp_dim, out_dim, stride=2)]
    layers += [residual(out_dim, out_dim) for _ in range(1, modules)]
    return nn.Sequential(*layers)

class model(hg_net):
 # 继承hg_net模块，就是把hg_net的所有定义拿过来可以直接调用，
 #这里的model就是CorNerNet model结构的所有
    def _pred_mod(self, dim):  # 用1*1的核升维或者降维到dim个channel
        return nn.Sequential(
            convolution(3, 256, 256, with_bn=False),
            nn.Conv2d(256, dim, (1, 1)) 
        )

    def _merge_mod(self):
        return nn.Sequential(
            nn.Conv2d(256, 256, (1, 1), bias=False), 
            #用1*1的核升维或者降维到256个channel
            nn.BatchNorm2d(256)
        )

    def __init__(self):
        stacks  = 2   #堆叠的沙漏网络，2个沙漏堆一起
        pre     = nn.Sequential(
            convolution(7, 3, 128, stride=2),
            residual(128, 256, stride=2)
        )
 #传入一个（B，N，W，H），B是batch，N是channel，W,H是feature map的维度，
#进入上面的pre模块 “预热”了一下下，让channel的数量变为256，
#这个就是为了下一步好和hg_net 模块好衔接，hg_net 模块从256开始増维到512，
#然后降维256（都是对feature map的channel操作，维度越大，map的size越小，
#维度小，size大，故称沙漏网络） 

        hg_mods = nn.ModuleList([
            hg_module(
                5, [256, 256, 384, 384, 384, 512], [2, 2, 2, 2, 2, 4],
                make_pool_layer=make_pool_layer,
                make_hg_layer=make_hg_layer
            ) for _ in range(stacks)  #2   2次重堆叠沙漏
        ])
        cnvs    = nn.ModuleList([convolution(3, 256, 256) for _ in range(stacks)])
        inters  = nn.ModuleList([residual(256, 256) for _ in range(stacks - 1)])
        cnvs_   = nn.ModuleList([self._merge_mod() for _ in range(stacks - 1)])
        inters_ = nn.ModuleList([self._merge_mod() for _ in range(stacks - 1)])
#  convs ,inters, cnvs_,inters_,这几个都是把一些操作组合成ModuleList执行

        hgs = hg(pre, hg_mods, cnvs, inters, cnvs_, inters_) 
        #沙漏网络的构建，hg_mods里面有2个沙漏网络，强调一下
        tl_modules = nn.ModuleList([corner_pool(256, TopPool, LeftPool) for _ in range(stacks)])
        br_modules = nn.ModuleList([corner_pool(256, BottomPool, RightPool) for _ in range(stacks)])

        tl_heats = nn.ModuleList([self._pred_mod(80) for _ in range(stacks)])
        br_heats = nn.ModuleList([self._pred_mod(80) for _ in range(stacks)])
        #t1_modules, br_moudles,tl_heats,br_heats 层是抽取特征图的信息，构建pred