Faster-RCNN系列 一（训练部分代码，python）
Faster-RCNN系列 一（训练部分代码，python）
本文注意从两部分讲解训练部分：
一：mobilenetV2
二：ResNet50+fpn
检查设备信息，并将设备信息打印：
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
检查是否有保存权重的文件，如果没有马厩makedirs一个
if not os.path.exists("save_weights"):
os.makedirs("save_weights")
数据集的处理，使用transforms进行增强处理，此处只将训练数据进行了水平翻转的处理，并训练和验证数据进行了格式的转换
data_transform = {
"train": transforms.Compose([transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val": transforms.Compose([transforms.ToTensor()])
}
对VOC数据集进行路径操作，找到需要使用的文件路径，包括训练数据和验证数据的txt文件
VOC_root = "./"
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], "train.txt")
val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt")
将数据集进行加载，此处要传入几个参数，每一Batch的大小，工作的进程数，特别注意的是，因为图像是标注过的，所以包括两个信息，一个是图像本身信息，一个是targets信息，所以在转换的时候要将targets对应
@staticmethod
def collate_fn(batch):
return tuple(zip(*batch))
train_data_loader = torch.utils.data.DataLoader(train_data_set,
batch_size=batch_size,
shuffle=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
val_data_set_loader = torch.utils.data.DataLoader(val_data_set,
batch_size=batch_size,
shuffle=False,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
数据定义完之后就是创建model，其中numclass为分类的类别个数+1
model = create_model(num_classes=21)
create_model部分：可以使用mobilenetV2网络，也可以使用ResNet50+fpn网络
def create_model(num_classes):
# https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features
backbone.out_channels = 1280
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def create_model(num_classes, device):
backbone = resnet50_fpn_backbone()
# 训练自己数据集时不要修改这里的91，修改的是传入的num_classes参数
model = FasterRCNN(backbone=backbone, num_classes=91)
# 载入预训练模型权重
# https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
weights_dict = torch.load("./backbone/fasterrcnn_resnet50_fpn_coco.pth", map_location=device)
missing_keys, unexpected_keys = model.load_state_dict(weights_dict, strict=False)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
将定义的model指定到设备上
model.to(device)
train_loss = []
learning_rate = []
val_mAP = []
此处如果是用预训练的backbone，现将backbone冻结，对Faster后半部分进行预训练
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)
num_epochs = 5
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
utils.train_one_epoch(model, optimizer, train_data_loader,
device, epoch, print_freq=50,
train_loss=train_loss, train_lr=learning_rate)
# evaluate on the test dataset
utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)
torch.save(model.state_dict(), "./save_weights/pretrain.pth")
再解冻之前训练的部分，对全部的网络结构进行训练，此时可以将backbone的底层权重冻结，对其余数据进行训练，并绘制loss和map曲线
# 冻结backbone部分底层权重
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
# 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)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5,
gamma=0.33)
num_epochs = 20
for epoch in range(num_epochs):
# train for one epoch, printing every 50 iterations
utils.train_one_epoch(model, optimizer, train_data_loader,
device, epoch, print_freq=50,
train_loss=train_loss, train_lr=learning_rate)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)
# save weights
if epoch > 10:
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))
# plot loss and lr curve
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)
# plot mAP curve
if len(val_mAP) != 0:
from plot_curve import plot_map
plot_map(val_mAP)
import os
import torch
import torchvision
from torchvision.ops import misc
import transforms
from network_files.faster_rcnn_framework import FasterRCNN
from network_files.rpn_function import AnchorsGenerator
from backbone.mobilenetv2_model import MobileNetV2
from my_dataset import VOC2012DataSet
from train_utils import train_eval_utils as utils
def create_model(num_classes):
# https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features
backbone.out_channels = 1280
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 检查保存权重文件夹是否存在，不存在则创建
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_root = "./"
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], "train.txt")
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = 8
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(train_data_set,
batch_size=batch_size,
shuffle=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# load validation data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt")
val_data_set_loader = torch.utils.data.DataLoader(val_data_set,
batch_size=batch_size,
shuffle=False,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# create model num_classes equal background + 20 classes
model = create_model(num_classes=21)
# print(model)
model.to(device)
train_loss = []
learning_rate = []
val_mAP = []
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# first frozen backbone and train 5 epochs #
# 首先冻结前置特征提取网络权重（backbone），训练rpn以及最终预测网络部分 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
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)
num_epochs = 5
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
utils.train_one_epoch(model, optimizer, train_data_loader,
device, epoch, print_freq=50,
train_loss=train_loss, train_lr=learning_rate)
# evaluate on the test dataset
utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)
torch.save(model.state_dict(), "./save_weights/pretrain.pth")
# # # # # # # # # # # # # # # # # # # # # # # # # # # #
# second unfrozen backbone and train all network #
# 解冻前置特征提取网络权重（backbone），接着训练整个网络权重 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # #
# 冻结backbone部分底层权重
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
# 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)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5,
gamma=0.33)
num_epochs = 20
for epoch in range(num_epochs):
# train for one epoch, printing every 50 iterations
utils.train_one_epoch(model, optimizer, train_data_loader,
device, epoch, print_freq=50,
train_loss=train_loss, train_lr=learning_rate)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)
# save weights
if epoch > 10:
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))
# plot loss and lr curve
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)
# plot mAP curve
if len(val_mAP) != 0:
from plot_curve import plot_map
plot_map(val_mAP)
# model.eval()
# x = [torch.rand(3, 300, 400), torch.rand(3, 400, 400)]
# predictions = model(x)
# print(predictions)
if __name__ == "__main__":
version = torch.version.__version__[:5] # example: 1.6.0
# 因为使用的官方的混合精度训练是1.6.0后才支持的，所以必须大于等于1.6.0
if version < "1.6.0":
raise EnvironmentError("pytorch version must be 1.6.0 or above")
main()
@霹雳吧啦Wz，感谢分享的视频，边学习，边复盘