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