def create_model(num_classes, device):
# 注意，这里的backbone默认使用的是FrozenBatchNorm2d，即不会去更新bn参数
# 目的是为了防止batch_size太小导致效果更差(如果显存很小，建议使用默认的FrozenBatchNorm2d)
# 如果GPU显存很大可以设置比较大的batch_size就可以将norm_layer设置为普通的BatchNorm2d
# trainable_layers包括['layer4', 'layer3', 'layer2', 'layer1', 'conv1']， 5代表全部训练
backbone = resnet50_fpn_backbone(norm_layer=torch.nn.BatchNorm2d,
trainable_layers=3)
# 训练自己数据集时不要修改这里的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
def create_model(num_classes):
# 如果训练的时候没有使用冻结BN，预测的时候也不要使用
# mobileNetv2+faster_RCNN
# backbone = MobileNetV2().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'],
# output_size=[7, 7],
# sampling_ratio=2)
#
# model = FasterRCNN(backbone=backbone,
# num_classes=num_classes,
# rpn_anchor_generator=anchor_generator,
# box_roi_pool=roi_pooler)
# resNet50+fpn+faster_RCNN
# 注意，这里的norm_layer要和训练脚本中保持一致
backbone = resnet50_fpn_backbone(norm_layer=torch.nn.BatchNorm2d)
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_score_thresh=0.5)
return model
def create_model(num_classes):
# 创建retinanet_res50_fpn模型
# skip P2 because it generates too many anchors (according to their paper)
# 注意，这里的backbone默认使用的是FrozenBatchNorm2d，即不会去更新bn参数
# 目的是为了防止batch_size太小导致效果更差(如果显存很小，建议使用默认的FrozenBatchNorm2d)
# 如果GPU显存很大可以设置比较大的batch_size就可以将norm_layer设置为普通的BatchNorm2d
backbone = resnet50_fpn_backbone(norm_layer=torch.nn.BatchNorm2d,
returned_layers=[2, 3, 4],
extra_blocks=LastLevelP6P7(256, 256),
trainable_layers=3)
model = RetinaNet(backbone, num_classes)
# 载入预训练权重
# https://download.pytorch.org/models/retinanet_resnet50_fpn_coco-eeacb38b.pth
weights_dict = torch.load("./backbone/retinanet_resnet50_fpn.pth",
map_location='cpu')
# 删除分类器部分的权重，因为自己的数据集类别与预训练数据集类别(91)不一定致，如果载入会出现冲突
del_keys = [
"head.classification_head.cls_logits.weight",
"head.classification_head.cls_logits.bias"
]
for k in del_keys:
del weights_dict[k]
print(model.load_state_dict(weights_dict, strict=False))
return model
def create_model(num_classes):
# resNet50+fpn+retinanet
# 注意，这里的norm_layer要和训练脚本中保持一致
backbone = resnet50_fpn_backbone(norm_layer=torch.nn.BatchNorm2d,
returned_layers=[2, 3, 4],
extra_blocks=LastLevelP6P7(256, 256))
model = RetinaNet(backbone, num_classes)
return model
def main(parser_data):
device = torch.device(
parser_data.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
data_transform = {"val": transforms.Compose([transforms.ToTensor()])}
# read class_indict
label_json_path = './pascal_voc_classes.json'
assert os.path.exists(
label_json_path), "json file {} dose not exist.".format(
label_json_path)
json_file = open(label_json_path, 'r')
class_dict = json.load(json_file)
category_index = {v: k for k, v in class_dict.items()}
VOC_root = parser_data.data_path
# 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))
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = parser_data.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
# load validation data set
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,
pin_memory=True,
collate_fn=val_data_set.collate_fn)
# create model num_classes equal background + 20 classes
# 注意，这里的norm_layer要和训练脚本中保持一致
backbone = resnet50_fpn_backbone(norm_layer=torch.nn.BatchNorm2d)
model = FasterRCNN(backbone=backbone,
num_classes=parser_data.num_classes + 1)
# 载入你自己训练好的模型权重
weights_path = parser_data.weights
assert os.path.exists(weights_path), "not found {} file.".format(
weights_path)
weights_dict = torch.load(weights_path, map_location=device)
model.load_state_dict(weights_dict['model'])
# print(model)
model.to(device)
# evaluate on the test dataset
coco = get_coco_api_from_dataset(val_data_set)
iou_types = ["bbox"]
coco_evaluator = CocoEvaluator(coco, iou_types)
cpu_device = torch.device("cpu")
model.eval()
with torch.no_grad():
for image, targets in tqdm(val_data_set_loader, desc="validation..."):
# 将图片传入指定设备device
image = list(img.to(device) for img in image)
# inference
outputs = model(image)
outputs = [{k: v.to(cpu_device)
for k, v in t.items()} for t in outputs]
res = {
target["image_id"].item(): output
for target, output in zip(targets, outputs)
}
coco_evaluator.update(res)
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
coco_eval = coco_evaluator.coco_eval["bbox"]
# calculate COCO info for all classes
coco_stats, print_coco = summarize(coco_eval)
# calculate voc info for every classes(IoU=0.5)
voc_map_info_list = []
for i in range(len(category_index)):
stats, _ = summarize(coco_eval, catId=i)
voc_map_info_list.append(" {:15}: {}".format(category_index[i + 1],
stats[1]))
print_voc = "\n".join(voc_map_info_list)
print(print_voc)
# 将验证结果保存至txt文件中
with open("record_mAP.txt", "w") as f:
record_lines = [
"COCO results:", print_coco, "", "mAP(IoU=0.5) for each category:",
print_voc
]
f.write("\n".join(record_lines))
