def init_weights(self, pretrained=None): if pretrained is None: return load_checkpoint(self, pretrained, strict=False) # we don't have bbox_head now print("init weight from {}".format(pretrained))
def init_weights(self, pretrained=None): if isinstance(pretrained, str): logger = logging.getLogger() load_checkpoint(self, pretrained, strict=False, logger=logger) elif pretrained is None: for m in self.modules(): if isinstance(m, nn.Conv2d): kaiming_init(m) elif isinstance(m, (_BatchNorm, nn.GroupNorm)): constant_init(m, 1) if self.dcn is not None: for m in self.modules(): if isinstance(m, Bottleneck) and hasattr( m, "conv2_offset"): constant_init(m.conv2_offset, 0) if self.zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): constant_init(m.norm3, 0) elif isinstance(m, BasicBlock): constant_init(m.norm2, 0) else: raise TypeError("pretrained must be a str or None")
def initialize_model(args): global model, voxel_generator cfg = Config.fromfile(args.config) model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) if args.checkpoint is not None: load_checkpoint(model, args.checkpoint, map_location="cpu") # print(model) if args.fp16: print("cast model to fp16") model = model.half() model = model.cuda() model.eval() global device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") range = cfg.voxel_generator.range voxel_size = cfg.voxel_generator.voxel_size max_points_in_voxel = cfg.voxel_generator.max_points_in_voxel max_voxel_num = cfg.voxel_generator.max_voxel_num[1] voxel_generator = VoxelGenerator(voxel_size=voxel_size, point_cloud_range=range, max_num_points=max_points_in_voxel, max_voxels=max_voxel_num) return model
def init_weights(self, pretrained=None): if pretrained is None: return try: load_checkpoint(self, pretrained, strict=False) print("init weight from {}".format(pretrained)) except: print("no pretrained model at {}".format(pretrained))
def main(): # Usage: python visualize.py <config_path> <checkpoint_path> parser = argparse.ArgumentParser() parser.add_argument("config_path", type=str, help="Path to the configuration file") parser.add_argument("checkpoint_path", type=str, help="Path to the model checkpoint file") args = parser.parse_args() cfg = torchie.Config.fromfile(args.config_path) cfg.data.val.test_mode = True # build the dataloader dataset = build_dataset(cfg.data.val) data_loader = build_dataloader( dataset, batch_size=1, workers_per_gpu=1, dist=False, shuffle=False, ) # build the model and load checkpoint model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) checkpoint = load_checkpoint(model, args.checkpoint_path, map_location="cpu") model = MegDataParallel(model, device_ids=[0]) device = torch.device("cuda") visualize(model, data_loader, device)
def init_weights(self, pretrained=None): if isinstance(pretrained, str): logger = logging.getLogger() load_checkpoint(self, pretrained, strict=False, logger=logger) elif pretrained is None: for m in self.features.modules(): if isinstance(m, nn.Conv2d): kaiming_init(m) elif isinstance(m, nn.BatchNorm2d): constant_init(m, 1) elif isinstance(m, nn.Linear): normal_init(m, std=0.01) else: raise TypeError("pretrained must be a str or None") for m in self.extra.modules(): if isinstance(m, nn.Conv2d): xavier_init(m, distribution="uniform") constant_init(self.l2_norm, self.l2_norm.scale)
def main(): # config cfg = Config.fromfile(CONFIG_FILE) # model loading model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) checkpoint = load_checkpoint(model, CHECK_POINT, map_location="cpu") model = model.cuda() model.eval() # data loader dataset = build_dataset(cfg.data.test) data_loader = build_dataloader( dataset, batch_size=cfg.data.samples_per_gpu, workers_per_gpu=cfg.data.workers_per_gpu, dist=False, shuffle=False, ) # infer detections = [] for i, data_batch in enumerate(data_loader): print("step:", i) with torch.no_grad(): outputs = batch_processor( model, data_batch, train_mode=False, local_rank=0, ) for output in outputs: token = output["metadata"]["token"] for k, v in output.items(): if k not in [ "metadata", ]: output[k] = v.to(cpu_device) detections.update({ token: output, }) all_predictions = all_gather(detections)
def main(): cfg = Config.fromfile( 'configs/nusc/pp/nusc_centerpoint_pp_02voxel_two_pfn_10sweep_demo.py') model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) dataset = build_dataset(cfg.data.val) data_loader = DataLoader( dataset, batch_size=1, sampler=None, shuffle=False, num_workers=8, collate_fn=collate_kitti, pin_memory=False, ) checkpoint = load_checkpoint( model, 'work_dirs/centerpoint_pillar_512_demo/latest.pth', map_location="cpu") model.eval() model = model.cuda() cpu_device = torch.device("cpu") points_list = [] gt_annos = [] detections = [] for i, data_batch in enumerate(data_loader): info = dataset._nusc_infos[i] gt_annos.append(convert_box(info)) points = data_batch['points'][:, 1:4].cpu().numpy() with torch.no_grad(): outputs = batch_processor( model, data_batch, train_mode=False, local_rank=0, ) for output in outputs: for k, v in output.items(): if k not in [ "metadata", ]: output[k] = v.to(cpu_device) detections.append(output) points_list.append(points.T) print( 'Done model inference. Please wait a minute, the matplotlib is a little slow...' ) for i in range(len(points_list)): visual(points_list[i], gt_annos[i], detections[i], i) print("Rendered Image {}".format(i)) image_folder = 'demo' video_name = 'video.avi' images = [img for img in os.listdir(image_folder) if img.endswith(".png")] images.sort(key=lambda img_name: int(img_name.split('.')[0][4:])) frame = cv2.imread(os.path.join(image_folder, images[0])) height, width, layers = frame.shape video = cv2.VideoWriter(video_name, 0, 1, (width, height)) cv2_images = [] for image in images: cv2_images.append(cv2.imread(os.path.join(image_folder, image))) for img in cv2_images: video.write(img) cv2.destroyAllWindows() video.release() print("Successfully save video in the main folder")
def main(): args = parse_args() assert args.out or args.show or args.json_out, ( "Please specify at least one operation (save or show the results) " 'with the argument "--out" or "--show" or "--json_out"' ) if args.out is not None and not args.out.endswith((".pkl", ".pickle")): raise ValueError("The output file must be a pkl file.") if args.json_out is not None and args.json_out.endswith(".json"): args.json_out = args.json_out[:-5] cfg = torchie.Config.fromfile(args.config) # set cudnn_benchmark if cfg.get("cudnn_benchmark", False): torch.backends.cudnn.benchmark = True # cfg.model.pretrained = None cfg.data.test.test_mode = True # cfg.data.val.test_mode = True # init distributed env first, since logger depends on the dist info. if args.launcher == "none": distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) # build the dataloader # TODO: support multiple images per gpu (only minor changes are needed) dataset = build_dataset(cfg.data.test) # dataset = build_dataset(cfg.data.val) data_loader = build_dataloader( dataset, batch_size=cfg.data.samples_per_gpu, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False, ) # build the model and load checkpoint model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) checkpoint = load_checkpoint(model, args.checkpoint, map_location="cpu") # old versions did not save class info in checkpoints, this walkaround is # for backward compatibility if "CLASSES" in checkpoint["meta"]: model.CLASSES = checkpoint["meta"]["CLASSES"] else: model.CLASSES = dataset.CLASSES model = MegDataParallel(model, device_ids=[0]) result_dict, detections = test( data_loader, model, save_dir=None, distributed=distributed ) for k, v in result_dict["results"].items(): print(f"Evaluation {k}: {v}") rank, _ = get_dist_info() if args.out and rank == 0: print("\nwriting results to {}".format(args.out)) torchie.dump(detections, args.out) if args.txt_result: res_dir = os.path.join(os.getcwd(), "predictions") for dt in detections: with open( os.path.join(res_dir, "%06d.txt" % int(dt["metadata"]["token"])), "w" ) as fout: lines = kitti.annos_to_kitti_label(dt) for line in lines: fout.write(line + "\n") ap_result_str, ap_dict = kitti_evaluate( "/data/Datasets/KITTI/Kitti/object/training/label_2", res_dir, label_split_file="/data/Datasets/KITTI/Kitti/ImageSets/val.txt", current_class=0, ) print(ap_result_str)
def main(): # torch.manual_seed(0) # torch.backends.cudnn.deterministic = True # torch.backends.cudnn.benchmark = False # np.random.seed(0) args = parse_args() cfg = Config.fromfile(args.config) # update configs according to CLI args if args.work_dir is not None: cfg.work_dir = args.work_dir distributed = torch.cuda.device_count() > 1 if distributed: if args.launcher == "pytorch": torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") cfg.local_rank = args.local_rank elif args.launcher == "slurm": proc_id = int(os.environ["SLURM_PROCID"]) ntasks = int(os.environ["SLURM_NTASKS"]) node_list = os.environ["SLURM_NODELIST"] num_gpus = torch.cuda.device_count() cfg.gpus = num_gpus torch.cuda.set_device(proc_id % num_gpus) addr = subprocess.getoutput( f"scontrol show hostname {node_list} | head -n1") # specify master port port = None if port is not None: os.environ["MASTER_PORT"] = str(port) elif "MASTER_PORT" in os.environ: pass # use MASTER_PORT in the environment variable else: # 29500 is torch.distributed default port os.environ["MASTER_PORT"] = "29501" # use MASTER_ADDR in the environment variable if it already exists if "MASTER_ADDR" not in os.environ: os.environ["MASTER_ADDR"] = addr os.environ["WORLD_SIZE"] = str(ntasks) os.environ["LOCAL_RANK"] = str(proc_id % num_gpus) os.environ["RANK"] = str(proc_id) dist.init_process_group(backend="nccl") cfg.local_rank = int(os.environ["LOCAL_RANK"]) cfg.gpus = torch.distributed.get_world_size() else: cfg.gpus = args.gpus # init logger before other steps logger = get_root_logger(cfg.log_level) logger.info("Distributed testing: {}".format(distributed)) logger.info(f"torch.backends.cudnn.benchmark: {torch.backends.cudnn.benchmark}") model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) if args.testset: print("Use Test Set") dataset = build_dataset(cfg.data.test) else: print("Use Val Set") dataset = build_dataset(cfg.data.val) data_loader = build_dataloader( dataset, batch_size=cfg.data.samples_per_gpu if not args.speed_test else 1, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False, ) checkpoint = load_checkpoint(model, args.checkpoint, map_location="cpu") # put model on gpus if distributed: # model = apex.parallel.convert_syncbn_model(model) model = DistributedDataParallel( model.cuda(cfg.local_rank), device_ids=[cfg.local_rank], output_device=cfg.local_rank, # broadcast_buffers=False, find_unused_parameters=True, ) else: # model = fuse_bn_recursively(model) model = model.cuda() model.eval() mode = "val" prog_bar = None logger.info(f"work dir: {args.work_dir}") if cfg.local_rank == 0: prog_bar = torchie.ProgressBar(len(data_loader.dataset) // cfg.gpus) detections = {} cpu_device = torch.device("cpu") start = time.time() start = int(len(dataset) / 3) end = int(len(dataset) * 2 /3) time_start = 0 time_end = 0 for i, data_batch in enumerate(data_loader): if i == start: torch.cuda.synchronize() time_start = time.time() if i == end: torch.cuda.synchronize() time_end = time.time() with torch.no_grad(): outputs = batch_processor( model, data_batch, train_mode=False, local_rank=args.local_rank, ) for output in outputs: token = output["metadata"]["token"] for k, v in output.items(): if k not in [ "metadata", ]: output[k] = v.to(cpu_device) detections.update( {token: output,} ) if args.local_rank == 0: if prog_bar is not None: prog_bar.update() synchronize() all_predictions = all_gather(detections) print("\n Total time per frame: ", (time_end - time_start) / (end - start)) if args.local_rank != 0: return predictions = {} for p in all_predictions: predictions.update(p) if not os.path.exists(args.work_dir): os.makedirs(args.work_dir) save_pred(predictions, args.work_dir) result_dict, _ = dataset.evaluation(copy.deepcopy(predictions), output_dir=args.work_dir, testset=args.testset) if result_dict is not None: for k, v in result_dict["results"].items(): print(f"Evaluation {k}: {v}") if args.txt_result: assert False, "No longer support kitti"
def main(): # torch.manual_seed(0) # torch.backends.cudnn.deterministic = True # torch.backends.cudnn.benchmark = False # np.random.seed(0) args = parse_args() cfg = Config.fromfile(args.config) cfg.local_rank = args.local_rank # update configs according to CLI args if args.work_dir is not None: cfg.work_dir = args.work_dir distributed = False if "WORLD_SIZE" in os.environ: distributed = int(os.environ["WORLD_SIZE"]) > 1 if distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") cfg.gpus = torch.distributed.get_world_size() else: cfg.gpus = args.gpus # init logger before other steps logger = get_root_logger(cfg.log_level) logger.info("Distributed testing: {}".format(distributed)) logger.info( f"torch.backends.cudnn.benchmark: {torch.backends.cudnn.benchmark}") model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) dataset = build_dataset(cfg.data.val) data_loader = build_dataloader( dataset, batch_size=cfg.data.samples_per_gpu, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False, ) checkpoint = load_checkpoint(model, args.checkpoint, map_location="cpu") # put model on gpus if distributed: model = apex.parallel.convert_syncbn_model(model) model = DistributedDataParallel( model.cuda(cfg.local_rank), device_ids=[cfg.local_rank], output_device=cfg.local_rank, # broadcast_buffers=False, find_unused_parameters=True, ) else: model = model.cuda() model.eval() mode = "val" logger.info(f"work dir: {args.work_dir}") if cfg.local_rank == 0: prog_bar = torchie.ProgressBar(len(data_loader.dataset) // cfg.gpus) detections = {} cpu_device = torch.device("cpu") for i, data_batch in enumerate(data_loader): with torch.no_grad(): outputs = batch_processor( model, data_batch, train_mode=False, local_rank=args.local_rank, ) for output in outputs: token = output["metadata"]["token"] for k, v in output.items(): if k not in [ "metadata", ]: output[k] = v.to(cpu_device) detections.update({ token: output, }) if args.local_rank == 0: prog_bar.update() synchronize() all_predictions = all_gather(detections) if args.local_rank != 0: return predictions = {} for p in all_predictions: predictions.update(p) result_dict, _ = dataset.evaluation(predictions, output_dir=args.work_dir) for k, v in result_dict["results"].items(): print(f"Evaluation {k}: {v}") if args.txt_result: res_dir = os.path.join(os.getcwd(), "predictions") for k, dt in predictions.items(): with open( os.path.join(res_dir, "%06d.txt" % int(dt["metadata"]["token"])), "w") as fout: lines = kitti.annos_to_kitti_label(dt) for line in lines: fout.write(line + "\n") ap_result_str, ap_dict = kitti_evaluate( "/data/Datasets/KITTI/Kitti/object/training/label_2", res_dir, label_split_file="/data/Datasets/KITTI/Kitti/ImageSets/val.txt", current_class=0, ) print(ap_result_str)
def main(): # torch.manual_seed(0) # torch.backends.cudnn.deterministic = True # torch.backends.cudnn.benchmark = False # np.random.seed(0) args = parse_args() cfg = Config.fromfile(args.config) cfg.local_rank = args.local_rank # update configs according to CLI args if args.work_dir is not None: cfg.work_dir = args.work_dir distributed = False if "WORLD_SIZE" in os.environ: distributed = int(os.environ["WORLD_SIZE"]) > 1 if distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") cfg.gpus = torch.distributed.get_world_size() else: cfg.gpus = args.gpus # init logger before other steps logger = get_root_logger(cfg.log_level) logger.info("Distributed testing: {}".format(distributed)) logger.info( f"torch.backends.cudnn.benchmark: {torch.backends.cudnn.benchmark}") model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) if args.testset: print("Use Test Set") dataset = build_dataset(cfg.data.test) else: print("Use Val Set") dataset = build_dataset(cfg.data.val) data_loader = build_dataloader( dataset, batch_size=cfg.data.samples_per_gpu if not args.speed_test else 1, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False, ) checkpoint = load_checkpoint(model, args.checkpoint, map_location="cpu") # put model on gpus if distributed: model = apex.parallel.convert_syncbn_model(model) model = DistributedDataParallel( model.cuda(cfg.local_rank), device_ids=[cfg.local_rank], output_device=cfg.local_rank, # broadcast_buffers=False, find_unused_parameters=True, ) else: # model = fuse_bn_recursively(model) model = model.cuda() model.eval() mode = "val" logger.info(f"work dir: {args.work_dir}") if cfg.local_rank == 0: prog_bar = torchie.ProgressBar(len(data_loader.dataset) // cfg.gpus) detections = {} cpu_device = torch.device("cpu") start = time.time() start = int(len(dataset) / 3) end = int(len(dataset) * 2 / 3) time_start = 0 time_end = 0 for i, data_batch in enumerate(data_loader): if i == start: torch.cuda.synchronize() time_start = time.time() if i == end: torch.cuda.synchronize() time_end = time.time() with torch.no_grad(): outputs = batch_processor( model, data_batch, train_mode=False, local_rank=args.local_rank, ) for output in outputs: token = output["metadata"]["token"] for k, v in output.items(): if k not in [ "metadata", ]: output[k] = v.to(cpu_device) detections.update({ token: output, }) if args.local_rank == 0: prog_bar.update() synchronize() all_predictions = all_gather(detections) print("\n Total time per frame: ", (time_end - time_start) / (end - start)) if args.local_rank != 0: return predictions = {} for p in all_predictions: predictions.update(p) if not os.path.exists(args.work_dir): os.makedirs(args.work_dir) save_pred(predictions, args.work_dir) with open(os.path.join(args.work_dir, 'prediction.pkl'), 'rb') as f: predictions = pickle.load(f) result_dict, _ = dataset.evaluation(copy.deepcopy(predictions), output_dir=args.work_dir, testset=args.testset) if result_dict is not None: for k, v in result_dict["results"].items(): print(f"Evaluation {k}: {v}") if args.txt_result: assert False, "No longer support kitti"
def main(): # torch.manual_seed(0) # torch.backends.cudnn.deterministic = True # torch.backends.cudnn.benchmark = False # np.random.seed(0) args = parse_args() cfg = Config.fromfile(args.config) cfg.local_rank = args.local_rank # update configs according to CLI args if args.work_dir is not None: cfg.work_dir = args.work_dir distributed = False if "WORLD_SIZE" in os.environ: distributed = int(os.environ["WORLD_SIZE"]) > 1 if distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") cfg.gpus = torch.distributed.get_world_size() else: cfg.gpus = args.gpus # init logger before other steps logger = get_root_logger(cfg.log_level) logger.info("Distributed testing: {}".format(distributed)) logger.info( f"torch.backends.cudnn.benchmark: {torch.backends.cudnn.benchmark}") torch.cuda.empty_cache() model = build_detector(cfg.nohead_model, train_cfg=None, test_cfg=cfg.test_cfg) if args.testset: print("Use Test Set") dataset = build_dataset(cfg.data.test) else: print("Use Val Set") dataset = build_dataset(cfg.data.val) data_loader = build_dataloader( dataset, batch_size=cfg.data.samples_per_gpu if not args.speed_test else 1, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False, ) checkpoint = load_checkpoint(model, args.checkpoint, map_location="cpu") # put model on gpus # model = fuse_bn_recursively(model) model = model.cuda() model.eval() mode = "val" logger.info(f"work dir: {args.work_dir}") if cfg.local_rank == 0: prog_bar = torchie.ProgressBar(len(data_loader.dataset) // cfg.gpus) detections = {} cpu_device = torch.device("cpu") start = time.time() start = int(len(dataset) / 3) end = int(len(dataset) * 2 / 3) time_start = 0 time_end = 0 device = torch.device(args.local_rank) POINTS_NUM = 2 for i, data_batch in enumerate(data_loader): if i == start: torch.cuda.synchronize() time_start = time.time() if i == end: torch.cuda.synchronize() time_end = time.time() with torch.no_grad(): sample = example_to_device(data_batch, device=device) for i in range(len(sample["metadata"])): sample["metadata"][i]['image_prefix'] = None del sample["metadata"] del sample["points"] #del sample["shape"] sample["shape"] = torch.tensor(sample["shape"]) sample["voxels"] = sample["voxels"][0:POINTS_NUM, :, :] sample["num_points"] = sample["num_points"][0:POINTS_NUM] sample["coordinates"] = sample["coordinates"][0:POINTS_NUM, :] outputs = model(sample, return_loss=False) #outputs = batch_processor( # model, data_batch, train_mode=False, local_rank=args.local_rank, #) for k, t in sample.items(): print("====", k) print(t.shape) print("============== start =============") register_custom_op_symbolic("spconv::get_indice_pairs_3d", symbolic_get_indice_pairs_3d, 11) torch.onnx.export( model, # model being run sample, # model input (or a tuple for multiple inputs) "/workspace/data/center_point.onnx", # where to save the model (can be a file or file-like object) export_params= True, # store the trained parameter weights inside the model file opset_version=11, # the ONNX version to export the model to do_constant_folding=True ) # whether to execute constant folding for optimization print("============== finish =============") break for output in outputs: token = output["metadata"]["token"] for k, v in output.items(): if k not in [ "metadata", ]: output[k] = v.to(cpu_device) detections.update({ token: output, }) if args.local_rank == 0: prog_bar.update() synchronize() all_predictions = all_gather(detections) print("\n Total time per frame: ", (time_end - time_start) / (end - start)) if args.local_rank != 0: return predictions = {} for p in all_predictions: predictions.update(p) if not os.path.exists(args.work_dir): os.makedirs(args.work_dir) save_pred(predictions, args.work_dir) result_dict, _ = dataset.evaluation(copy.deepcopy(predictions), output_dir=args.work_dir, testset=args.testset) if result_dict is not None: for k, v in result_dict["results"].items(): print(f"Evaluation {k}: {v}") if args.txt_result: assert False, "No longer support kitti"
def main(): args = parse_args() print(args) assert args.out or args.show or args.json_out, ( 'Please specify at least one operation (save or show the results) with the argument "--out" or "--show" or "--json_out"' ) if args.out is not None and not args.out.endswith((".pkl", ".pickle")): raise ValueError("The output file must be a pkl file.") if args.json_out is not None and args.json_out.endswith(".json"): args.json_out = args.json_out[:-5] cfg = torchie.Config.fromfile(args.config) if cfg.get("cudnn_benchmark", False): # False torch.backends.cudnn.benchmark = True # cfg.model.pretrained = None # cfg.data.test.test_mode = True cfg.data.val.test_mode = True # init distributed env first, since logger depends on the dist info. if args.launcher == "none": distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) # build the dataloader, TODO: support multiple images per gpu (only minor changes are needed) dataset = build_dataset(cfg.data.val) batch_size = cfg.data.samples_per_gpu num_workers = cfg.data.workers_per_gpu data_loader = DataLoader( dataset, batch_size=batch_size, sampler=None, num_workers=num_workers, collate_fn=collate_kitti, shuffle=False, ) # build the model and load checkpoint model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) checkpoint_path = os.path.join(cfg.work_dir, args.checkpoint) checkpoint = load_checkpoint(model, checkpoint_path, map_location="cpu") # old versions did not save class info in checkpoints, this walkaround is for backward compatibility if "CLASSES" in checkpoint["meta"]: model.CLASSES = checkpoint["meta"]["CLASSES"] else: model.CLASSES = dataset.CLASSES model = MegDataParallel(model, device_ids=[0]) if args.eval_id is None: result_dict, detections = test(data_loader, model, save_dir=None, distributed=distributed) # for k, v in result_dict["results"].items(): # print(f"Evaluation {k}: {v}") for k, v in result_dict["results"].items(): print(f"Evaluation {k}: {v}") # f.write(f"\nEvaluation {k}: {v}\n") for k, v in result_dict["results_2"].items(): print(f"Evaluation {k}: {v}") # f.write(f"\nEvaluation {k}: {v}\n") # save mAP results to out.pkl file. # rank, _ = get_dist_info() # if args.out and rank == 0: # print("\nwriting results to {}".format(args.out)) # torchie.dump(detections, os.path.join(cfg.work_dir, args.out)) # if args.txt_result: # True # res_dir = os.path.join(cfg.work_dir, "predictions") # os.makedirs(res_dir, exist_ok=True) # for dt in detections: # with open(os.path.join(res_dir, "%06d.txt" % int(dt["metadata"]["token"])), "w") as fout: # lines = kitti.annos_to_kitti_label(dt) # for line in lines: # fout.write(line + "\n") # gt_labels_dir = data_root + "/KITTI/object/training/label_2" # label_split_file = data_root + "/KITTI/ImageSets/val.txt" # # todo: this evaluation is different from previous one # ap_result_str, ap_dict = kitti_evaluate(gt_labels_dir, res_dir, label_split_file=label_split_file, current_class=0,) # print(ap_result_str) else: assert type(args.eval_id) is list test_v2(data_loader, model, distributed=distributed, eval_id=args.eval_id, vis_id=args.vis_id)
def main(): cfg = Config.fromfile( 'configs/nusc/pp/nusc_centerpoint_pp_02voxel_two_pfn_10sweep_demo_export_onnx.py' ) model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) dataset = build_dataset(cfg.data.val) data_loader = DataLoader( dataset, batch_size=1, sampler=None, shuffle=False, num_workers=8, collate_fn=collate_kitti, pin_memory=False, ) checkpoint = load_checkpoint(model, './latest.pth', map_location="cpu") model.eval() model = model.cuda() gpu_device = torch.device("cuda") points_list = [] gt_annos = [] detections = [] data_iter = iter(data_loader) data_batch = next(data_iter) pp_model = PointPillars(model) points = data_batch['points'][:, 1:4].cpu().numpy() with torch.no_grad(): example = example_to_device(data_batch, gpu_device, non_blocking=False) example["voxels"] = torch.zeros( (example["voxels"].shape[0], example["voxels"].shape[1], 10), dtype=torch.float32, device=gpu_device) example.pop("metadata") example.pop("points") example["shape"] = torch.tensor(example["shape"], dtype=torch.int32, device=gpu_device) model(example) torch.onnx.export( model.reader, (example["voxels"], example["num_voxels"], example["coordinates"]), "onnx_model/pfe.onnx", opset_version=11) rpn_input = torch.zeros((1, 64, 512, 512), dtype=torch.float32, device=gpu_device) torch.onnx.export(pp_model, rpn_input, "onnx_model/rpn.onnx", opset_version=11) print("Done")