def main():
args = arguments.parse_args()
dataset = VOCDataSet("data", args.dataset, enabled_flip=False)
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=5,
collate_fn=collate_fn)
print("[+] Bootstrapping model")
print("[+] Loading rpn model from %s" % args.rpn)
print("[+] Loading fast rcnn model from %s" % args.fastrcnn)
fast_rcnn_checkpoint = torch.load(args.fastrcnn)
rpn_checkpoint = torch.load(args.rpn)
fast_rcnn = FasterRCNN().cuda()
fast_rcnn.load_state_dict(fast_rcnn_checkpoint["state_dict"])
fast_rcnn.eval()
rpn = RegionProposalNetwork().cuda()
rpn.load_state_dict(rpn_checkpoint["state_dict"])
print("[+] Calculating Average Precision")
ap = get_ap(rpn,
fast_rcnn,
dataset,
dataloader,
name="%s_%s" % (args.dataset, args.name))
print("Average Precision=%.3f" % ap)
def get_model(num_classes, device="cuda:0"):
# 加载backbone时不需要加载官方模型
backbone = resnet_fpn_backbone('resnet50',
pretrained=False,
trainable_layers=3)
model = FasterRCNN(backbone, num_classes=91) # 加载官方模型，这里num_classes不能修改
# 加载fasterrcnn_resnet50_fpn官方模型
# https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
weights_dict = torch.load(
"C:/Users/lixiao/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
map_location=device)
missing_keys, unexpected_keys = model.load_state_dict(weights_dict,
strict=False)
overwrite_eps(model, 0.0)
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) # 修改predictor
return model
def main():
global args
args = arguments.parse_args()
experiment_env = create_experiment_dir()
assert args.rois_per_batch % args.batch_size == 0, "Uneven number of rois per image"
rois_per_image = args.rois_per_batch / args.batch_size
train_data = VOCDataSetROIs("data",
"train",
rois_per_image,
enabled_flip=True)
dataloader = DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
collate_fn=collate_rois_fn)
################### MODEL BOOTSRAP #####################
print("[+] Bootstrapping model")
if args.stage_2_path is not None:
print("[+] Loading stage 2 weights")
net = FasterRCNN(args.stage_2_path).cuda()
net.train()
if args.resume is not None:
print("[+] Resuming from %s" % args.resume)
checkpoint = torch.load(args.resume)
net.load_state_dict(checkpoint["state_dict"])
cross_entropy = CrossEntropyLoss(size_average=True).cuda()
smooth_l1_loss = SmoothL1Loss(size_average=False).cuda()
optimizer = opt.SGD(
[params for params in net.parameters() if params.requires_grad],
lr=args.lr,
momentum=args.momentum,
weight_decay=0.0005)
################### MODEL TRAINING #####################
print("[+] Training model")
start_epoch = 0 if args.resume is None else checkpoint["epoch"]
for epoch in xrange(start_epoch, args.epoch):
adjust_learning_rate(optimizer, epoch)
train(net, cross_entropy, smooth_l1_loss, optimizer, dataloader,
experiment_env, epoch)
def get_model(num_classes, device="cuda:0"):
# # https://download.pytorch.org/models/vgg16-397923af.pth
# # 如果使用mobilenetv2的话就下载对应预训练权重并注释下面三行，接着把mobilenetv2模型对应的两行代码注释取消掉
# vgg_feature = vgg(model_name="vgg16", weights_path="./backbone/vgg16.pth").features
# backbone = torch.nn.Sequential(*list(vgg_feature._modules.values())[:-1]) # 删除feature中最后的maxpool层
# backbone.out_channels = 512
# # https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
# # backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features
# # backbone.out_channels = 1280 # 设置对应backbone输出特征矩阵的channels
# anchor_generator = AnchorGenerator(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)
# 加载backbone时不需要加载官方模型
backbone = resnet_fpn_backbone('resnet50',
pretrained=False,
trainable_layers=3)
model = FasterRCNN(backbone, num_classes=91) # 加载官方模型，这里num_classes不能修改
# 加载fasterrcnn_resnet50_fpn官方模型
# https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
weights_dict = torch.load(
"C:/Users/lixiao/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
map_location=device)
missing_keys, unexpected_keys = model.load_state_dict(weights_dict,
strict=False)
overwrite_eps(model, 0.0)
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) # 修改predictor
return model