def create_model(num_classes=21, device=torch.device('cpu')):
# https://download.pytorch.org/models/resnet50-19c8e357.pth
# pre_train_path = "./src/resnet50.pth"
backbone = Backbone()
model = SSD300(backbone=backbone, num_classes=num_classes)
# https://ngc.nvidia.com/catalog/models -> search ssd -> download FP32
pre_ssd_path = "./src/nvidia_ssdpyt_fp32.pt"
if os.path.exists(pre_ssd_path) is False:
raise FileNotFoundError(
"nvidia_ssdpyt_fp32.pt not find in {}".format(pre_ssd_path))
pre_model_dict = torch.load(pre_ssd_path, map_location=device)
pre_weights_dict = pre_model_dict["model"]
# 删除类别预测器权重，注意，回归预测器的权重可以重用，因为不涉及num_classes
del_conf_loc_dict = {}
for k, v in pre_weights_dict.items():
split_key = k.split(".")
if "conf" in split_key:
continue
del_conf_loc_dict.update({k: v})
missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict,
strict=False)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
return model.to(device)
def create_model(num_classes, device=torch.device('cpu')):
# https://download.pytorch.org/models/resnet50-19c8e357.pth
# pre_train_path = "./src/resnet50.pth"
backbone = Backbone(pretrain_path=None)
model = SSD300(backbone=backbone, num_classes=num_classes)
pre_ssd_path = "./src/nvidia_ssdpyt_fp32.pt"
pre_model_dict = torch.load(pre_ssd_path, map_location=device)
pre_weights_dict = pre_model_dict["model"]
# 删除类别预测器权重，注意，回归预测器的权重可以重用，因为不涉及num_classes
del_conf_loc_dict = {}
for k, v in pre_weights_dict.items():
split_key = k.split(".")
if "conf" in split_key:
continue
del_conf_loc_dict.update({k: v})
missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict,
strict=False)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
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.Resize(),
transforms.ToTensor(),
transforms.Normalization()])
}
# 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)
json_file.close()
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
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_dataset = VOCDataSet(VOC_root, "2012", transforms=data_transform["val"], train_set="val.txt")
val_dataset_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=nw,
pin_memory=True,
collate_fn=val_dataset.collate_fn)
# create model num_classes equal background + 20 classes
backbone = Backbone()
model = SSD300(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)
model.load_state_dict(torch.load(weights_path, map_location=device)['model'])
# print(model)
model.to(device)
# evaluate on the test dataset
coco = get_coco_api_from_dataset(val_dataset)
iou_types = ["bbox"]
coco_evaluator = CocoEvaluator(coco, iou_types)
cpu_device = torch.device("cpu")
model.eval()
with torch.no_grad():
for images, targets in tqdm(val_dataset_loader, desc="validation..."):
# 将图片传入指定设备device
images = torch.stack(images, dim=0).to(device)
# inference
results = model(images)
outputs = []
for index, (bboxes_out, labels_out, scores_out) in enumerate(results):
# 将box的相对坐标信息（0-1）转为绝对值坐标(xmin, ymin, xmax, ymax)
height_width = targets[index]["height_width"]
# 还原回原图尺度
bboxes_out[:, [0, 2]] = bboxes_out[:, [0, 2]] * height_width[1]
bboxes_out[:, [1, 3]] = bboxes_out[:, [1, 3]] * height_width[0]
info = {"boxes": bboxes_out.to(cpu_device),
"labels": labels_out.to(cpu_device),
"scores": scores_out.to(cpu_device)}
outputs.append(info)
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))
def create_model(num_classes):
backbone = Backbone()
model = SSD300(backbone=backbone, num_classes=num_classes)
return model