Faster-RCNN代码解读4：辅助文件解读

Faster-RCNN代码解读4：辅助文件解读

前言

​ 因为最近打算尝试一下Faster-RCNN的复现，不要多想，我还没有厉害到可以一个人复现所有代码。所以，是参考别人的代码，进行自己的解读。

​ 代码来自于B站的UP主（大佬666），其把代码都放到了GitHub上了，我把链接都放到下面了（应该不算侵权吧，毕竟代码都开源了^_）：

b站链接：https://www.bilibili.com/video/BV1of4y1m7nj/?vd_source=afeab8b555e5eb1bfa1e7f267262cbf2GitHub链接：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

目的

​ 其实UP主已经做了很好的视频讲解了他的代码，只是有时候我还是喜欢阅读博客来学习，另外视频很长，6个小时，我看的时候容易睡着^_，所以才打算写博客记录一下学习笔记。

目前完成的内容

​ 第一篇：VOC数据集详细介绍

​ 第二篇：Faster-RCNN代码解读2：快速上手使用

​ 第三篇：Faster-RCNN代码解读3：制作自己的数据加载器

​ 第四篇：Faster-RCNN代码解读4：辅助文件解读（本文）

目录结构

1. 前言：

​ 本篇主要介绍的文件有：

split_data.py
plot_curve.py
draw_box_utils.py
train_mobilenetv2.py
backbone文件下的文件

​ 这些文件除去train_mobilenetv2.py外都是一些辅助的文件，读懂它们可以帮助我们后面理解Faster-RCNN的代码内容。

2. split_data.py文件：

​ 这个文件的主要作用：制作自己的train.txt和val.txt文件。其中，train.txt和val.txt文件就是VOC数据集ImageSets\\Main\\里的train.txt和val.txt文件，里面的数据分别是训练集和测试集的图片名字。

​ 其实这个文件作用并不大，因为数据集已经为我们提供了训练集和测试集的划分。不过，如果你用的自己的数据集，就可以用到它，它可以给你一个参考思路。

​ 下面，进行解读：

​ 首先，指定数据集的路径并验证该路径是否存在：

# 指定数据集地址
files_path = "./VOCdevkit/VOC2012/Annotations"
assert os.path.exists(files_path), "path: '{}' does not exist.".format(files_path)

​ 然后，设置训练集和测试集的比例：

# 设置验证集比例
val_rate = 0.5

​ 接着，将文件前缀和后缀分开，这是因为VOC数据集一个特点就是一张图片的注释、图像等文件前缀都是相同的：

# 切割文件名： 2007_000027.xml ---- [2007_000027,xml]，即获得2007_000027文件名字
files_name = sorted([file.split(".")[0] for file in os.listdir(files_path)])
# 获取总数
files_num = len(files_name)

​ 然后，使用random.sample随机采取所需的验证图片，该函数返回验证图片的索引，并迭代将训练图片和验证图片分开保存：

# 随机采取指定比例的数据，获取索引，并放入不同的列表中
val_index = random.sample(range(0, files_num), k=int(files_num*val_rate))
train_files = []
val_files = []
# 将上面采集的放入对应的列表中
for index, file_name in enumerate(files_name):# 如果索引在验证集的索引集合中if index in val_index:# 加入验证列表val_files.append(file_name)else:# 否则，加入训练集列表中train_files.append(file_name)

​ 最后，将上述内容保存到文件中即可：

# 将之保存到文件中
try:train_f = open("train.txt", "x")eval_f = open("val.txt", "x")train_f.write("\\n".join(train_files))eval_f.write("\\n".join(val_files))
except FileExistsError as e:print(e)exit(1)

3. plot_curve.py文件：

​ 这个文件是画图文件，其画的是损失函数、学习率和mAP图像。

​ 其实，这个文件很简单，主要的代码都是涉及matplotlib库的使用，所以不需要多说什么，可以看我写的代码注释：

def plot_loss_and_lr(train_loss, learning_rate):try:# 根据长度设置x轴的值x = list(range(len(train_loss)))fig, ax1 = plt.subplots(1, 1)   # 创建画布，注意只有一个ax1.plot(x, train_loss, 'r', label='loss')  # 画损失函数图# 美化图像ax1.set_xlabel("step")ax1.set_ylabel("loss")ax1.set_title("Train Loss and lr")plt.legend(loc='best')ax2 = ax1.twinx() # 启用右坐标轴ax2.plot(x, learning_rate, label='lr')  # 画学习率图ax2.set_ylabel("learning rate")ax2.set_xlim(0, len(train_loss))  # 设置横坐标整数间隔plt.legend(loc='best')handles1, labels1 = ax1.get_legend_handles_labels()  # 返回图例的句柄和标签，比如 legend为 loss,那么l就为losshandles2, labels2 = ax2.get_legend_handles_labels()plt.legend(handles1 + handles2, labels1 + labels2, loc='upper right')fig.subplots_adjust(right=0.8)  # 防止出现保存图片显示不全的情况# 保存图像fig.savefig('./loss_and_lr{}.png'.format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S")))plt.close()print("successful save loss curve! ")except Exception as e:print(e)def plot_map(mAP):try:# 根据长度设置x轴的值x = list(range(len(mAP)))plt.plot(x, mAP, label='mAp')   # 画mAP图# 美化图像plt.xlabel('epoch')plt.ylabel('mAP')plt.title('Eval mAP')plt.xlim(0, len(mAP))plt.legend(loc='best')# 保存plt.savefig('./mAP.png')plt.close()print("successful save mAP curve!")except Exception as e:print(e)

​ 我这里说一下上面涉及到的，而我们一般又不使用的代码。

• ax1.twinx()：作用是启用右坐标轴。

​ 如果你用过origin或者对画图了解一点，应该知道右坐标轴的意思（见下图），只是这个函数你可能没见过，其作用就是启用右坐标轴并返回一个操作对象。

• ax1.get_legend_handles_labels()：作用是返回这个坐标轴的图例句柄（即操作对象）和标签（即图例内容）

​ 这个说起来难以理解，但是举个例子就简单了。比如上图右上角的第一条曲线就是一个图例，其内容/标签值就是第一条曲线。而句柄就是操作对象，如果返回的对象为空，表示这个图没有使用图例。

4. draw_box_utils.py文件：

​ 这个文件的主要作用就是画出图像的边界框、类别信息和mask信息。

​ 首先，看最下面的函数draw_objs：

draw_objs函数

​ 作用：画出所有对象的边界框和mask。

​ 输入的参数：

​ 首先，过滤掉哪些概率值较低的边界框：

# 过滤掉低概率的目标
idxs = np.greater(scores, box_thresh)
# 需要同时处理boxes、classes、scores、masks
boxes = boxes[idxs]
classes = classes[idxs]
scores = scores[idxs]
if masks is not None:masks = masks[idxs]

​ 接着，判断过滤后，是否全部过滤掉，如果全部过滤掉就不需要画了：

# 如果boxes长度为0，表示所有的框都过滤了，就不需要画了
if len(boxes) == 0:return image

​ 然后，随机从定义的颜色列表中抽取颜色，生成一个待使用的颜色列表：

# 从定义的颜色列表中抽取颜色 
# ImageColor.getrgb 获取颜色的rgb值
colors = [ImageColor.getrgb(STANDARD_COLORS[cls % len(STANDARD_COLORS)]) for cls in classes]

​ 接着，开始画边界框，看注释即可：

# 如果需要画边界框
if draw_boxes_on_image:# 创建画图对象draw = ImageDraw.Draw(image)# 开始迭代绘图，因为一张图不知一个对象，所以需要画出所有的框for box, cls, score, color in zip(boxes, classes, scores, colors):# 边界框的坐标left, top, right, bottom = box# 绘制目标边界框，顺时针画图draw.line([(left, top), (left, bottom), (right, bottom),(right, top), (left, top)], width=line_thickness, fill=color)# 绘制类别和概率信息draw_text(draw, box.tolist(), int(cls), float(score), category_index, color, font, font_size)

​ 最后，画mask：

if draw_masks_on_image and (masks is not None):# 画出所有的maskimage = draw_masks(image, masks, colors, mask_thresh)

​ 其中上面的draw_text和draw_masks是文件中的另外两个函数，下面进行讲解。

draw_text函数

​ 作用：将目标边界框和类别信息绘制到图片上，是draw_obj的辅助函数。

​ 输入参数：

​ 首先，由于需要画文字，即需要创建文字对象：

# 创建字体对象，如果创建失败（比如作者用的字体你没有），就使用默认的字体
try:font = ImageFont.truetype(font, font_size)
except IOError:font = ImageFont.load_default()

​ 接下来，就是获取边界框的坐标信息并设置文字要显示在边界框的哪个位置：

# 获取坐标
left, top, right, bottom = box
# 将数字的类别转为真实的类别信息，并加上概率值构成“ person 99% ”这样的字符串
display_str = f"{category_index[str(cls)]}: {int(100 * score)}%"
# 设置字体的高度
display_str_heights = [font.getsize(ds)[1] for ds in display_str]
display_str_height = (1 + 2 * 0.05) * max(display_str_heights)# 如果文字的高度没有超过图像最高点
if top > display_str_height:# 设置文字的坐标text_top = top - display_str_heighttext_bottom = top
else:# 如果超过了，就设置文字的坐标为边界框的下面text_top = bottomtext_bottom = bottom + display_str_height

​ 最后，就是将边界框和文字画在图像上：

 # 开始画
for ds in display_str:# 获取文字的宽和高text_width, text_height = font.getsize(ds)margin = np.ceil(0.05 * text_width)# 画一个矩形draw.rectangle([(left, text_top),(left + text_width + 2 * margin, text_bottom)], fill=color)# 画文字draw.text((left + margin, text_top),ds,fill='black',font=font)left += text_width

draw_masks函数：

​ 这个函数比较简单，看注释即可：

def draw_masks(image, masks, colors, thresh: float = 0.7, alpha: float = 0.5):# 将图像转为array值np_image = np.array(image)# 过滤下maskmasks = np.where(masks > thresh, True, False)# colors = np.array(colors)img_to_draw = np.copy(np_image)# TODO: There might be a way to vectorize this# 将mask区域改变颜色for mask, color in zip(masks, colors):img_to_draw[mask] = colorout = np_image * (1 - alpha) + img_to_draw * alpha# 最后，将array转为图像return fromarray(out.astype(np.uint8))

5. backbone文件夹下的文件：

​ 这个文件夹下的内容就是骨干CNN架构的内容。其下有四个主要的文件：

resnet50+fpn
vgg
mobilenetv2
feature-pyramid-network

​ 这四个文件，其实没有什么好说的，因为都是根据网络架构来实现，对于我们来说，并不是很重要。当然，如果你感兴趣，可以在网络找到对应的架构图，然后参考代码自己实现，都是可以的。

6. train_mobilenetv2.py文件：

​ 这个文件和train_res50_fpn.py内容上都是相同的，代码大体也相似，作用就是训练backbone文件夹下的CNN架构。这里我以train_mobilenetv2.py来解读一下具体的内容。

main函数

​ 首先，肯定是指定一些参数变量，比如指定GPU、指定是否存在权重保存文件夹、指定预处理方、指定采用的数据集、batch_size大小等等：

# 指定GPU设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))# 用来保存coco_info的文件
# coco_info文件：
results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists("save_weights"):os.makedirs("save_weights")# 指定数据增强方式，即随机水平翻转（框和图片都要翻转）
data_transform = {"train": transforms.Compose([transforms.ToTensor(),transforms.RandomHorizontalFlip(0.5)]),"val": transforms.Compose([transforms.ToTensor()])
}# 指定VOC数据集地址----需要修改
VOC_root = "./"  # VOCdevkit
aspect_ratio_group_factor = 3
batch_size = 8 # batch size大小
amp = False  # 是否使用混合精度训练，需要GPU支持# 检查VOC数据集是否存在，否则报错
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))

​ 接着，定义数据集和数据集的加载器（看注释）：

# 加载数据集，使用我们自己定义的加载器来加载
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_dataset = VOCDataSet(VOC_root, "2012", data_transform["train"], "train.txt")
train_sampler = None# 是否按图片相似高宽比采样图片组成batch
# 使用的话能够减小训练时所需GPU显存，默认使用
if aspect_ratio_group_factor >= 0:train_sampler = torch.utils.data.RandomSampler(train_dataset)# 统计所有图像高宽比例在bins区间中的位置索引group_ids = create_aspect_ratio_groups(train_dataset,k=aspect_ratio_group_factor)# 每个batch图片从同一高宽比例区间中取train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, batch_size)# 使用多少个线程去加载图片
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])
print('Using %g dataloader workers' % nw)# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合batch
if train_sampler:# 如果按照图片高宽比采样图片，dataloader中需要使用batch_samplertrain_data_loader = torch.utils.data.DataLoader(train_dataset,batch_sampler=train_batch_sampler, # 与sampler类似，但是一次只返回一个batch的indices（索引），需要注意的是，一旦指定了这个参数，那么batch_size,shuffle,sampler,drop_last就不能再制定了（互斥——Mutually exclusive）pin_memory=True, # 如果设置为True，那么data loader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中num_workers=nw,  # 多线程读取数据collate_fn=train_dataset.collate_fn)  #  将一个list的sample组成一个mini-batch的函数
else:train_data_loader = torch.utils.data.DataLoader(train_dataset,batch_size=batch_size,shuffle=True,pin_memory=True,num_workers=nw,collate_fn=train_dataset.collate_fn)# 加载验证集数据集
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_dataset = VOCDataSet(VOC_root, "2012", data_transform["val"], "val.txt")
val_data_loader = torch.utils.data.DataLoader(val_dataset,batch_size=1,shuffle=False,pin_memory=True,num_workers=nw,collate_fn=val_dataset.collate_fn)

​ 然后，将定义模型并将模型放入GPU中，顺带定义一些变量：（注意，此时的模型为Faster-RCNN模型，只是里面的CNN架构为mobilenetv2）

# 创建模型，类别为固定的20+一个背景
model = create_model(num_classes=21)
# print(model)# 放入GPU中
model.to(device)# 梯度缩放，即有些梯度很小，计算机无法存储完，就会下溢，这时将梯度放大，即可存储下来
scaler = torch.cuda.amp.GradScaler() if amp else None# 定义一些变量，主要用于后面的画图
train_loss = [] # 训练损失
learning_rate = []  # 学习率
val_map = []    # 验证集的mAP值

​ 下面就是重头戏了，就是训练网络。作者这里采取如下的训练思路：首先，冻结CNN架构的权重（即让这部分不求梯度），然后用于训练RPN网络，这一阶段只训练5个epoch。然后，解冻CNN架构的权重，开始训练整个网络，这里需要注意，作者认为CNN架构最前面的几层是公用，又加上数据不多，因此前面几层不参与训练，即将其冻结权重。

​ 有了这个思路，看代码就很简单了，具体的可以看注释内容：

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#  first frozen backbone and train 5 epochs                   #
#  首先冻结前置特征提取网络权重（backbone），训练rpn以及最终预测网络部分 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ## 不求backbone的梯度，即不调节它们
for param in model.backbone.parameters():param.requires_grad = False# define optimizer
# 确定要优化的参数
params = [p for p in model.parameters() if p.requires_grad]
# 定义优化器
optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005)# 在前5个epoch训练对后面的参数微调
init_epochs = 5
for epoch in range(init_epochs):# train for one epoch, printing every 10 iterationsmean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,device, epoch, print_freq=50,warmup=True, scaler=scaler)train_loss.append(mean_loss.item())learning_rate.append(lr)# 在测试集上验证coco_info = utils.evaluate(model, val_data_loader, device=device)# 训练信息写入文件with open(results_file, "a") as f:# 写入的数据包括coco指标还有loss和learning rateresult_info = [f"{i:.4f}" for i in coco_info + [mean_loss.item()]] + [f"{lr:.6f}"]txt = "epoch:{} {}".format(epoch, '  '.join(result_info))f.write(txt + "\\n")val_map.append(coco_info[1])  # pascal mAP# 保存权重
torch.save(model.state_dict(), "./save_weights/pretrain.pth")# # # # # # # # # # # # # # # # # # # # # # # # # # # #
#  second unfrozen backbone and train all network     #
#  解冻前置特征提取网络权重（backbone），接着训练整个网络权重  #
# # # # # # # # # # # # # # # # # # # # # # # # # # # ## 冻结backbone部分底层权重：认为前面几层是公用的特征+data很少，训练整个网络不够，因此冻结部分层（官方实现方法）
for name, parameter in model.backbone.named_parameters():split_name = name.split(".")[0]if split_name in ["0", "1", "2", "3"]:parameter.requires_grad = False # 冻结else:parameter.requires_grad = True  # 解冻# 确定哪些参数需要训练
params = [p for p in model.parameters() if p.requires_grad]
# 定义优化器
optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005)
# 学习调整方法，每三次调整一次，乘以0.33
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.33)
# 开始训练，调整参数
num_epochs = 20
for epoch in range(init_epochs, num_epochs+init_epochs, 1):# 开始训练一个epoch，每50次迭代打印依次损失值mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,device, epoch, print_freq=50,warmup=True, scaler=scaler)# 保存平均损失和当前学习率train_loss.append(mean_loss.item())learning_rate.append(lr)# 更新学习率lr_scheduler.step()# 在测试集上验证coco_info = utils.evaluate(model, val_data_loader, device=device)# 将训练信息写入文件with open(results_file, "a") as f:# 写入的数据包括coco指标还有loss和learning rateresult_info = [f"{i:.4f}" for i in coco_info + [mean_loss.item()]] + [f"{lr:.6f}"]txt = "epoch:{} {}".format(epoch, '  '.join(result_info))f.write(txt + "\\n")val_map.append(coco_info[1])  # pascal mAP# 仅保存最后5个epoch的权重# 还需要保存一些优化器、学习率等的参数if epoch in range(num_epochs+init_epochs)[-5:]:save_files = {'model': model.state_dict(),'optimizer': optimizer.state_dict(),'lr_scheduler': lr_scheduler.state_dict(),'epoch': epoch}torch.save(save_files, "./save_weights/mobile-model-{}.pth".format(epoch))

​ 最后，就是画出损失函数、学习率和mAP图像即可：

# 画损失函数图和学习率图
if len(train_loss) != 0 and len(learning_rate) != 0:from plot_curve import plot_loss_and_lrplot_loss_and_lr(train_loss, learning_rate)# 画mAP图
if len(val_map) != 0:from plot_curve import plot_mapplot_map(val_map)

create_model函数：

​ 这个函数是上面main函数中的创建模型函数。

​ 这个函数里面涉及到了很多其它的函数，我会在下一篇进行讲解，这里仅仅做概述：

def create_model(num_classes):# https://download.pytorch.org/models/vgg16-397923af.pth# 如果使用vgg16的话就下载对应预训练权重并取消下面注释，接着把mobilenetv2模型对应的两行代码注释掉# vgg_feature = vgg(model_name="vgg16", weights_path="./backbone/vgg16.pth").features# backbone = torch.nn.Sequential(*list(vgg_feature._modules.values())[:-1])  # 删除features中最后一个Maxpool层# backbone.out_channels = 512# 拥有预训练权重# https://download.pytorch.org/models/mobilenet_v2-b0353104.pth# backbone采用mobilenetv2backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features# 设置输出数目backbone.out_channels = 1280  # 设置对应backbone输出特征矩阵的channels# 生成anchor# size即尺寸# aspect_ratios即缩放因子anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512),),aspect_ratios=((0.5, 1.0, 2.0),))# 生成roiroi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],  # 在哪些特征层上进行roi poolingoutput_size=[7, 7],   # roi_pooling输出特征矩阵尺寸sampling_ratio=2)  # 采样率# 创建Faster-RCNN模型，并指定相关参数model = FasterRCNN(backbone=backbone,num_classes=num_classes,rpn_anchor_generator=anchor_generator,box_roi_pool=roi_pooler)return model

​ 这里需要注意的是AnchorsGenerator中的size参数有五个值，而在原论文中只有三个值128, 256, 512，加上三个缩放因子，就生成3*3=9个anchors。这里有五个值，是作者改变的，可能是方便小目标的检测。

7. 总结：

​ 上面解读的文件中，最重要的就是train_mobilenetv2.py文件，这个文件就是训练Faster-RCNN的文件，只是采用的CNN架构为mobilenetV2。

​ 另外，值得一提的是：如果后期你要调试代码，查看变量的值，就需要运行该文件。