上次记录了下 TensorFlow 版本，这次由于剪枝需要，尝试下 PyTorch 版本。
源码目录总览
yolov3-ultralytics
├── cfg // 网络定义文件
│ ├── yolov3.cfg
│ ├── yolov3-spp.cfg
│ ├── yolov3-tiny.cfg
├── data // 数据配置
│ ├── samples // 示例图片，detect.py 检测的就是这里的图片
│ ├── coco.names // coco 用于检测的 80 个类别的名字
│ ├── coco_paper.names // coco 原始 91 个类别的名字
│ ├── coco2014.data // coco 2014 版本的训练测试路径配置
│ └── coco2017.data // coco 2017 版本的训练测试路径配置
├── utils // 核心代码所在文件夹
│ ├── __init__.py
│ ├── adabound.py
│ ├── datasets.py
│ ├── google_utils.py
│ ├── layers.py
│ ├── parse_config.py
│ ├── torch_utils.py
│ └── utils.py
├── weights // 模型所在路径
│ ├── yolov3-spp-ultralytics.weights // 原始 YOLOV3 模型格式
│ └── yolov3-spp-ultralytics.pt // PyTorch 模型格式
├── detect.py // demo 代码
├── models.py // 核心代码
├── test.py // 测试数据集 mAP
├── train.py // 模型训练
└── tutorial.ipynb // 使用教程接下来，我们从数据加载、网络定义、网络训练、mAP 测试等角度来仔细过一遍代码
Inference
入手的第一步是，运行作者提供的 detect.py 这个 demo 检测示例。代码的运行参数熟悉检测的应该很容易从字面上理解，需要强调的是 --source 这个参数如果是 '0' 则会调用摄像头模型，默认是会读取 'data/samples' 文件夹下的示例图片。
整个代码按照 网络初始化 -> 模型加载 -> 输入图片加载 -> 前向推理 -> NMS 后处理 -> 可视化检测结果 的顺序来执行的。
Initialize Model
model = Darknet(opt.cfg, imgsz)这里会根据网络定义文件(本文以 'cfg/yolov3-spp.cfg' 为例) 来定义网络结构，这个和官方 Darknet 是保存一样的设计的，这也是为什么这份代码生成的 pytorch 模型能与官方 Darknet 的模型能够无缝转换的原因。
这里贴一张 'cfg/yolov3-spp.cfg' 网络结构示意图：
具体实现参考 models.py 这份代码。这里有个有意思的地方是默认会使用 thop 按照固定输入(1, 3, 480, 640)尺度统计参数量和计算量，看来在 pytorch 中 thop 的使用率还是挺高的。
imgsz 这个参数目的是为了 onnx 准备的，因为要固定输入尺寸。
Dataloader
dataset = LoadImages(source, img_size=imgsz)这里的数据预处理包括 padding_resize 和 BGR2EGB，具体参考 utils/datasets.py。推理前图片还要经过 img /= 255.0 归一化操作，就是不知道为什么作者不把这个操作放到 datasets.py 里了。
Inference
除了常规的前向推理外，代码中还定义了一种 Augment images 推理方式，有兴趣的可以去瞅瞅。主要关注的是 YOLOLayer 层的推理，这基本上和 TensorFlow 版本 的差不多 。
YOLOV3 有三个尺度的预测输出，以上面图示 512x512 尺度的输入，coco 数据集训练的模型的第一个尺度（最小 feature map) 输出为例。YOLOLayer 前一层的输出维度为 16 x 16 x 255，16 x 16 对应的是 feature map size，255=3*（4+1+80)，其中 3 是由于每个尺度上设计了 3 种 Anchor, 4 分别是中心点xy的偏移、宽高的偏移量，1 代表是否包含目标，80 则是具体类别的置信度。
流程上作者将 YOLOLayer 前一层的输出维度进行变换:
# p.view(bs, 255, 16, 16) -- > (bs, 3, 16, 16, 85) # (bs, anchors, grid, grid, xywh + classes)
p = p.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction随后按照论文中的方式对预测结果进行 decode:
io = p.clone() # inference output
io[..., :2] = torch.sigmoid(io[..., :2]) + self.grid # xy
io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh # wh yolo method
io[..., :4] *= self.stride
torch.sigmoid_(io[..., 4:]) # 会改变 io
return io.view(bs, -1, self.no), p # view [1, 3, 16, 16, 85] as [1, 768, 85]具体实现参考 models.py
Train
训练这块算是这个版本的亮点之处了，相比于官方 Darknet 改进点还是很多的。
Dataloader
# Dataset
dataset = LoadImagesAndLabels(train_path, img_size, batch_size,
augment=True,
hyp=hyp, # augmentation hyperparameters
rect=opt.rect, # rectangular training
cache_images=opt.cache_images,
single_cls=opt.single_cls)
# Dataloader
batch_size = min(batch_size, len(dataset))
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
num_workers=nw,
shuffle=not opt.rect, # Shuffle=True unless rectangular training is used
pin_memory=True,
collate_fn=dataset.collate_fn)区别于推理环节只需要加载和处理图片，这里加载的是类似于 COO 的数据集。以训练 coco2017.data 为例，训练数据 train_path 保存的是训练集的图片路径，LoadImagesAndLabels 中会对应加载 txt 格式的标注文件。
以 train2017/000000391895.jpg 为例，标注文件格式如下 (cx, cy, w, h)
3 0.6490546875000001 0.7026527777777778 0.17570312500000002 0.59325
0 0.65128125 0.47923611111111114 0.2404375 0.8353611111111112
0 0.765 0.5468611111111111 0.05612500000000001 0.13361111111111112
1 0.7833203125 0.5577777777777778 0.047859375 0.09716666666666667该图片尺寸(HxW)为 360x640, 训练时如果设置的训练尺度为 512，那么这张图片需要先保持高宽比缩放到 288x512 然后再 padding 到 512x512，当然对应的 label 也要有所调整，最后图片要 BGR2EGB 转换。
数据增广包括：
具体参考 utils/datasets.py。
考虑到存在 crop 形式的数据增广，这里区别于 TensorFlow 版本 并没有在数据层就把 Anchor 匹配给做了，而是在训练中做了。
Optimizer
比较有特点的是，代码中将可训练的参数划分为三组，卷积层权重为一组，bias 为一组，其它的参数为一组
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in dict(model.named_parameters()).items():
if '.bias' in k:
pg2 += [v] # biases
elif 'Conv2d.weight' in k:
pg1 += [v] # apply weight_decay
else:
pg0 += [v] # all else训练中，这三组参数设置不同学习速率
if opt.adam:
# hyp['lr0'] *= 0.1 # reduce lr (i.e. SGD=5E-3, Adam=5E-4)
optimizer = optim.Adam(pg0, lr=hyp['lr0'])
# optimizer = AdaBound(pg0, lr=hyp['lr0'], final_lr=0.1)
else:
# momentum 一次设置后面都默认是这么多
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)基础学习率采用 cosine
lf = lambda x: (((1 + math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.95 + 0.05 # cosine
同时，代码中也设计了 warmup 训练方式：
nb = len(dataloader) # number of batches
n_burn = max(3 * nb, 500)
# Burn-in
if ni <= n_burn:
xi = [0, n_burn] # x interp
model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(1, np.interp(ni, xi, [1, 64 / batch_size]).round())
for j, x in enumerate(optimizer.param_groups): # pg0 pg1 pg2
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
x['weight_decay'] = np.interp(ni, xi, [0.0, hyp['weight_decay'] if j == 1 else 0.0])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']])| learning rate | weight decay |
 |
 |
| momentum | accumulate |
 |
 |
|
|
 |
值得注意的是：Optimize 是每隔 accumulate 个 batch 才更新模型的，也就是说不考虑 warmup 的情况下，训练的实际 batch_size = 64。
accumulate = max(round(64 / batch_size), 1)Others
同其它版本一样，这份代码也支持多尺度训练，且是每 effective bs = batch_size * accumulate 更新一次训练尺度
训练还采用了 Model Exponential Moving Average 机制
ema = torch_utils.ModelEMA(model)除了一些预定义的超参数外，还要关注下 model.gr 和 model.class_weights，同预定义的 cls, cls_pw, obj, obj_pw 一样这些都是计算 loss 所准备的一些权重项
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weightsLOSS
loss, loss_items = compute_loss(pred, targets, model)这部分定义在 utils/utils.py 文件中。
首先将 Anchor 和 GT 尝试进行匹配，匹配方式是中心点方式匹配，这里存在一个隐含的风险是有可能存在 GT 没有任何匹配 Anchor，其他版本中是会为该 GT 寻找一个最大匹配项，这里没有，默认舍弃了。
匹配函数 build_targets：
def build_targets(p, targets, model):
# Build targets for compute_loss(), input targets(image,class,x,y,w,h)
nt = targets.shape[0]
tcls, tbox, indices, anch = [], [], [], []
gain = torch.ones(6, device=targets.device) # normalized to gridspace gain
off = torch.tensor([[1, 0], [0, 1], [-1, 0], [0, -1]], device=targets.device).float() # overlap offsets
style = None
multi_gpu = type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
for i, j in enumerate(model.yolo_layers): # 每个尺度输出
anchors = model.module.module_list[j].anchor_vec if multi_gpu else model.module_list[j].anchor_vec
gain[2:] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
na = anchors.shape[0] # number of anchors
at = torch.arange(na).view(na, 1).repeat(1, nt) # anchor tensor, same as .repeat_interleave(nt)
# Match targets to anchors
a, t, offsets = [], targets * gain, 0
if nt:
# r = t[None, :, 4:6] / anchors[:, None] # wh ratio
# j = torch.max(r, 1. / r).max(2)[0] < model.hyp['anchor_t'] # compare
j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n) = wh_iou(anchors(3,2), gwh(n,2))
a, t = at[j], t.repeat(na, 1, 1)[j] # filter
# overlaps
gxy = t[:, 2:4] # grid xy
z = torch.zeros_like(gxy)
if style == 'rect2':
g = 0.2 # offset
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
a, t = torch.cat((a, a[j], a[k]), 0), torch.cat((t, t[j], t[k]), 0)
offsets = torch.cat((z, z[j] + off[0], z[k] + off[1]), 0) * g
elif style == 'rect4':
g = 0.5 # offset
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
l, m = ((gxy % 1. > (1 - g)) & (gxy < (gain[[2, 3]] - 1.))).T
a, t = torch.cat((a, a[j], a[k], a[l], a[m]), 0), torch.cat((t, t[j], t[k], t[l], t[m]), 0)
offsets = torch.cat((z, z[j] + off[0], z[k] + off[1], z[l] + off[2], z[m] + off[3]), 0) * g
# Define
b, c = t[:, :2].long().T # image, class
gxy = t[:, 2:4] # grid xy
gwh = t[:, 4:6] # grid wh
gij = (gxy - offsets).long()
gi, gj = gij.T # grid xy indices
# Append
indices.append((b, a, gj, gi)) # image, anchor, grid indices
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
anch.append(anchors[a]) # anchors
tcls.append(c) # class
if c.shape[0]: # if any targets
assert c.max() < model.nc, 'Model accepts %g classes labeled from 0-%g, however you labelled a class %g. ' \
'See https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data' % (
model.nc, model.nc - 1, c.max())
return tcls, tbox, indices, anch匹配函数返回四个变量 tcls, tbox, indices, anchors，还是以上面的示例 GT 同第一个尺度的 Anchor 匹配为例：
Anchor 预定义如下，标记部分是第一个尺度(最小 featuremap 用来检测大目标)：
anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,3261. 将这个尺度的 Anchor 映射到该 featuremap 尺度得到 anchors
2. 将 GT 映射到该 featuremap 尺度得到 t
3. 中心点方式匹配 ($Anchor_{num} x GT_{num}$)
这样的话匹配的 Anchor 索引 a=[0, 0,1,1, 2] 和 GT 为 下图 t，b 和 c 则分别代表匹配的 GT 在一个 batch 中的 index 和 类别 id：
b c
gij 计统计的是匹配的 GT 的中心点坐标，gi, gj 分别是 center_x 和 center_y
gi gj
最后输出的四个变量如下
tcls.append(c) # class
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
indices.append((b, a, gj, gi)) # image, anchor, grid indices
anch.append(anchors[a]) # anchors这里可以看出由 indices 和 anch 可以还原成匹配的 Anchor 所在的位置，而由 indices 和 tbox 可以还原出匹配的 GT 所在的位置
有了匹配结果我们就可以确定哪些作为正样本哪些作为负样本和每个正样本的回归目标了，继而可以计算出 分类和回归 Loss 了
是否包含目标 lobj 和目标类别分类 lcls 的 Loss 都采用 BCEWithLogitsLoss，如果超参数 fl_gamma 不为零，则采用 FocalLoss 加持，特别的 BCELoss 还分别有 cls_pw(1.0)、obj_pw(1.0) 两个权重加持。目标的回归 Loss 采用的 GIoU Loss，偶然在 CIoU 这篇论文开源的 github issue 里见到本代码作者询问 CIoU Loss 的一些实现细节，我想后期本代码作者应该也会加上 CIoU Loss 的吧。
除此之外，本代码作者还考虑将是否包含目标的 lobj Loss 用 giou ratio 加持，即 GIoU 越大，目标 label 越接近于 1(原本是固定值 1），目的应该是为了给匹配不好的 GT 弱一点的监督信号，降低训练难度。
tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * giou.detach().clamp(0).type(tobj.dtype) # giou ratio smooth
lobj += BCEobj(pi[..., 4], tobj) # obj loss最后，这三个 Loss 再分别由 giou(3.54)、obj(64.3) 、cls(37.4) 这三个权重值加持，至于这三个值是怎么来的，后面有时间我再研究下。
那么最后整个检测的 loss = lbox + lobj + lcls，如果你用过 TensorFlow 版本 的代码你会发现，相同的还是很多的，改进点就是这些权重项和 giou ratio 了。
回过头一看你会发现 model.class_weights 这一个权重项代码中并没有使用。
有关训练的部分就到此差不多，后面有时间再有坑填坑吧。
Evaluate COCO
这部分一般也都是工具性的调用，我就没仔细看了。反正你如果要对比不同的框架或者代码训练的模型的性能，选择其中一种然后按照对应的格式准备数据就好了。
Q
有关这份代码的简单理解就到此为止了，剩下一些要填的坑等后面有时间再补吧。待填坑包括：
1. 数据增广
2. 权重值设定的依据
3. 可以 试验下 HAMBox 中用回归后的 Anchor 替换初始 Anchor 来匹配的思想所加持的效果
4. 替换 GIoU Loss 为 CIoU Loss
接下来就是以此代码训练的模型，试试 剪枝 了
来源：oschina
链接：https://my.oschina.net/u/4385718/blog/4299622