Faster-RCNN代码解读4：辅助文件解读
前言
因为最近打算尝试一下Faster-RCNN的复现，不要多想，我还没有厉害到可以一个人复现所有代码。所以，是参考别人的代码，进行自己的解读。
代码来自于B站的UP主（大佬666），其把代码都放到了GitHub上了，我把链接都放到下面了（应该不算侵权吧，毕竟代码都开源了_）：
b站链接：https://www.bilibili.com/video/BV1of4y1m7nj/?vd_source=afeab8b555e5eb1bfa1e7f267262cbf2
GitHub链接：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing
目的
其实UP主已经做了很好的视频讲解了他的代码，只是有时候我还是喜欢阅读博客来学习，另外视频很长，6个小时，我看的时候容易睡着_，所以才打算写博客记录一下学习笔记。
目前完成的内容
第一篇：VOC数据集详细介绍
第二篇：Faster-RCNN代码解读2：快速上手使用
第三篇：Faster-RCNN代码解读3：制作自己的数据加载器
第四篇：Faster-RCNN代码解读4：辅助文件解读（本文）
目录结构
文章目录
- Faster-RCNN代码解读4：辅助文件解读
- 1. 前言：
- 2. split_data.py文件：
- 3. plot_curve.py文件：
- 4. draw_box_utils.py文件：
- 5. backbone文件夹下的文件：
- 6. train_mobilenetv2.py文件：
- 7. 总结：
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就为loss
handles2, 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。
输入的参数：
| 参数 | 意义 |
|---|---|
| image | 需要绘制的图片 |
| boxes | 目标边界框信息 |
| classes | 目标类别信息 |
| scores | 目标概率信息 |
| masks | 目标mask信息 |
| category_index | 类别与名称字典 |
| box_thresh | 过滤的概率阈值，默认为0.1 |
| mask_thresh | 同上，只是过滤的对象为mask，默认为0.5 |
| line_thickness | 边界框宽度 |
| font | 字体类型 |
| font_size | 字体大小 |
| draw_boxes_on_image | 是否将边界框画在图像上，默认为True |
| draw_masks_on_image | 是否将mask画在图像上，默认为Fasle |
首先，过滤掉哪些概率值较低的边界框：
# 过滤掉低概率的目标
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):
# 画出所有的mask
image = draw_masks(image, masks, colors, mask_thresh)
其中上面的draw_text和draw_masks是文件中的另外两个函数，下面进行讲解。
draw_text函数
作用：将目标边界框和类别信息绘制到图片上，是draw_obj的辅助函数。
输入参数：
| 参数 | 意义 |
|---|---|
| draw | 画图对象，可以使用画直线等等方法 |
| box | 一个边界框，里面有坐标信息 |
| cls | 对象的类别，为int值，需要使用category_index转为字符串值 |
| score | 对象的类别概率值 |
| category_index | 不同的索引对应的类别信息 |
| color | 使用的颜色 |
| font | 字体 |
| font_size | 字大小 |
首先，由于需要画文字，即需要创建文字对象：
# 创建字体对象，如果创建失败（比如作者用的字体你没有），就使用默认的字体
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_height
text_bottom = top
else:
# 如果超过了，就设置文字的坐标为边界框的下面
text_top = bottom
text_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)
# 过滤下mask
masks = 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] = color
out = 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_sampler
train_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 iterations
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)
# 在测试集上验证
coco_info = utils.evaluate(model, val_data_loader, device=device)
# 训练信息写入文件
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_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 rate
result_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_lr
plot_loss_and_lr(train_loss, learning_rate)
# 画mAP图
if len(val_map) != 0:
from plot_curve import plot_map
plot_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采用mobilenetv2
backbone = 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),))
# 生成roi
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_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。
另外，值得一提的是：如果后期你要调试代码，查看变量的值，就需要运行该文件。
