def inference(config_path, model_dir, result_path=None, predict_test=False, ckpt_path=None, ref_detfile=None): model_dir = pathlib.Path(model_dir) if predict_test: result_name = 'predict_test' else: result_name = 'eval_results' if result_path is None: result_path = model_dir / result_name else: result_path = pathlib.Path(result_path) config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config class_names = list(input_cfg.class_names) center_limit_range = model_cfg.post_center_limit_range ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) # net = second_builder.build(model_cfg, voxel_generator, target_assigner) batch_size = 1 num_worker = 1 net = second_builder.build(model_cfg, voxel_generator, target_assigner, 1) net.cuda() if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) if ckpt_path is None: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) else: torchplus.train.restore(ckpt_path, net) eval_dataset = input_reader_builder.build( input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=False, collate_fn=merge_second_batch) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 net.eval() result_path_step = result_path / f"step_{net.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) t = time.time() dt_annos = [] global_set = None print("Generate output labels...") bar = ProgressBar() bar.start(len(eval_dataset) // input_cfg.batch_size + 1) for example in iter(eval_dataloader): example = example_convert_to_torch(example, float_dtype) example_tuple = list(example.values()) batch_image_shape = example_tuple[8] example_tuple[8] = torch.from_numpy(example_tuple[8]) example_tuple[9] = torch.from_numpy(example_tuple[9]) # print("before", example) dt_annos = prediction_once( net, example_tuple, class_names, batch_image_shape, center_limit_range, model_cfg.lidar_input, global_set) return 0 bar.print_bar()
def evaluate(config_path, model_dir, result_path=None, predict_test=False, ckpt_path=None, ref_detfile=None, pickle_result=True): model_dir = pathlib.Path(model_dir) if predict_test: result_name = 'predict_test' else: result_name = 'eval_results' if result_path is None: result_path = model_dir / result_name else: result_path = pathlib.Path(result_path) config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config class_names = list(input_cfg.class_names) center_limit_range = model_cfg.post_center_limit_range ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) # net = second_builder.build(model_cfg, voxel_generator, target_assigner) net = second_builder.build(model_cfg, voxel_generator, target_assigner, input_cfg.batch_size) net.cuda() if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) if ckpt_path is None: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) else: torchplus.train.restore(ckpt_path, net) eval_dataset = input_reader_builder.build( input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=input_cfg.batch_size, shuffle=False, num_workers=input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 net.eval() result_path_step = result_path / f"step_{net.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) t = time.time() dt_annos = [] global_set = None print("Generate output labels...") bar = ProgressBar() bar.start(len(eval_dataset) // input_cfg.batch_size + 1) for example in iter(eval_dataloader): # evaluation example:[0:'voxels', 1:'num_points', 2:'coordinates', 3:'rect', # 4:'Trv2c', 5:'P2', # 6:'anchors', 7:'anchors_mask', 8:'image_idx', 9:'image_shape'] example = example_convert_to_torch(example, float_dtype) example_tuple = list(example.values()) example_tuple[8] = torch.from_numpy(example_tuple[8]) example_tuple[9] = torch.from_numpy(example_tuple[9]) if(example_tuple[6].size()[0] != input_cfg.batch_size): continue if pickle_result: dt_annos += predict_kitti_to_anno( net, example_tuple, class_names, center_limit_range, model_cfg.lidar_input, global_set) # print("shut train/py L703") # return 0 else: _predict_kitti_to_file(net, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input) bar.print_bar() sec_per_example = len(eval_dataset) / (time.time() - t) print(f'generate label finished({sec_per_example:.2f}/s). start eval:') print(f"avg forward time per example: {net.avg_forward_time:.3f}") print(f"avg postprocess time per example: {net.avg_postprocess_time:.3f}") if not predict_test: gt_annos = [info["annos"] for info in eval_dataset.dataset.kitti_infos] if(len(gt_annos)%2 != 0): del gt_annos[-1] if not pickle_result: dt_annos = kitti.get_label_annos(result_path_step) result = get_official_eval_result(gt_annos, dt_annos, class_names) print(result) result = get_coco_eval_result(gt_annos, dt_annos, class_names) print(result) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f)
def train(config_path, model_dir, result_path=None, create_folder=False, display_step=50, summary_step=5, pickle_result=True): """train a VoxelNet model specified by a config file. """ if create_folder: if pathlib.Path(model_dir).exists(): model_dir = torchplus.train.create_folder(model_dir) model_dir = pathlib.Path(model_dir) model_dir.mkdir(parents=True, exist_ok=True) eval_checkpoint_dir = model_dir / 'eval_checkpoints' eval_checkpoint_dir.mkdir(parents=True, exist_ok=True) if result_path is None: result_path = model_dir / 'results' config_file_bkp = "pipeline.config" config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) shutil.copyfile(config_path, str(model_dir / config_file_bkp)) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config class_names = list(input_cfg.class_names) ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) ###################### # BUILD TARGET ASSIGNER ###################### bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) ###################### # BUILD NET ###################### center_limit_range = model_cfg.post_center_limit_range # net = second_builder.build(model_cfg, voxel_generator, target_assigner) net = second_builder.build(model_cfg, voxel_generator, target_assigner, input_cfg.batch_size) net.cuda() # net_train = torch.nn.DataParallel(net).cuda() print("num_trainable parameters:", len(list(net.parameters()))) # for n, p in net.named_parameters(): # print(n, p.shape) ###################### # BUILD OPTIMIZER ###################### # we need global_step to create lr_scheduler, so restore net first. torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) gstep = net.get_global_step() - 1 optimizer_cfg = train_cfg.optimizer if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) optimizer = optimizer_builder.build(optimizer_cfg, net.parameters()) if train_cfg.enable_mixed_precision: loss_scale = train_cfg.loss_scale_factor mixed_optimizer = torchplus.train.MixedPrecisionWrapper( optimizer, loss_scale) else: mixed_optimizer = optimizer # must restore optimizer AFTER using MixedPrecisionWrapper torchplus.train.try_restore_latest_checkpoints(model_dir, [mixed_optimizer]) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, optimizer, gstep) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 ###################### # PREPARE INPUT ###################### dataset = input_reader_builder.build( input_cfg, model_cfg, training=True, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataset = input_reader_builder.build( eval_input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) def _worker_init_fn(worker_id): time_seed = np.array(time.time(), dtype=np.int32) np.random.seed(time_seed + worker_id) print(f"WORKER {worker_id} seed:", np.random.get_state()[1][0]) dataloader = torch.utils.data.DataLoader( dataset, batch_size=input_cfg.batch_size, shuffle=True, num_workers=input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch, worker_init_fn=_worker_init_fn) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=eval_input_cfg.batch_size, shuffle=False, num_workers=eval_input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch) data_iter = iter(dataloader) ###################### # TRAINING ###################### log_path = model_dir / 'log.txt' logf = open(log_path, 'a') logf.write(proto_str) logf.write("\n") summary_dir = model_dir / 'summary' summary_dir.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(str(summary_dir)) total_step_elapsed = 0 remain_steps = train_cfg.steps - net.get_global_step() t = time.time() ckpt_start_time = t total_loop = train_cfg.steps // train_cfg.steps_per_eval + 1 # total_loop = remain_steps // train_cfg.steps_per_eval + 1 clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch if train_cfg.steps % train_cfg.steps_per_eval == 0: total_loop -= 1 mixed_optimizer.zero_grad() try: for _ in range(total_loop): if total_step_elapsed + train_cfg.steps_per_eval > train_cfg.steps: steps = train_cfg.steps % train_cfg.steps_per_eval else: steps = train_cfg.steps_per_eval for step in range(steps): lr_scheduler.step() try: example = next(data_iter) except StopIteration: print("end epoch") if clear_metrics_every_epoch: net.clear_metrics() data_iter = iter(dataloader) example = next(data_iter) example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] example_tuple = list(example_torch.values()) example_tuple[11] = torch.from_numpy(example_tuple[11]) example_tuple[12] = torch.from_numpy(example_tuple[12]) assert 13==len(example_tuple), "something wring with training input size!" # training example:[0:'voxels', 1:'num_points', 2:'coordinates', 3:'rect', # 4:'Trv2c', 5:'P2', # 6:'anchors', 7:'anchors_mask', 8:'labels', 9:'reg_targets', 10:'reg_weights', # 11:'image_idx', 12:'image_shape'] # ret_dict = net(example_torch) # training input from example # print("example[0] size", example_tuple[0].size()) pillar_x = example_tuple[0][:,:,0].unsqueeze(0).unsqueeze(0) pillar_y = example_tuple[0][:,:,1].unsqueeze(0).unsqueeze(0) pillar_z = example_tuple[0][:,:,2].unsqueeze(0).unsqueeze(0) pillar_i = example_tuple[0][:,:,3].unsqueeze(0).unsqueeze(0) num_points_per_pillar = example_tuple[1].float().unsqueeze(0) # Find distance of x, y, and z from pillar center # assuming xyres_16.proto coors_x = example_tuple[2][:, 3].float() coors_y = example_tuple[2][:, 2].float() # self.x_offset = self.vx / 2 + pc_range[0] # self.y_offset = self.vy / 2 + pc_range[1] # this assumes xyres 20 # x_sub = coors_x.unsqueeze(1) * 0.16 + 0.1 # y_sub = coors_y.unsqueeze(1) * 0.16 + -39.9 # here assumes xyres 16 x_sub = coors_x.unsqueeze(1) * 0.16 + 0.08 y_sub = coors_y.unsqueeze(1) * 0.16 + -39.6 ones = torch.ones([1, 100],dtype=torch.float32, device=pillar_x.device ) x_sub_shaped = torch.mm(x_sub, ones).unsqueeze(0).unsqueeze(0) y_sub_shaped = torch.mm(y_sub, ones).unsqueeze(0).unsqueeze(0) num_points_for_a_pillar = pillar_x.size()[3] mask = get_paddings_indicator(num_points_per_pillar, num_points_for_a_pillar, axis=0) mask = mask.permute(0, 2, 1) mask = mask.unsqueeze(1) mask = mask.type_as(pillar_x) coors = example_tuple[2] anchors = example_tuple[6] labels = example_tuple[8] reg_targets = example_tuple[9] input = [pillar_x, pillar_y, pillar_z, pillar_i, num_points_per_pillar, x_sub_shaped, y_sub_shaped, mask, coors, anchors, labels, reg_targets] ret_dict = net(input) assert 10==len(ret_dict), "something wring with training output size!" # return 0 # ret_dict { # 0:"loss": loss, # 1:"cls_loss": cls_loss, # 2:"loc_loss": loc_loss, # 3:"cls_pos_loss": cls_pos_loss, # 4:"cls_neg_loss": cls_neg_loss, # 5:"cls_preds": cls_preds, # 6:"dir_loss_reduced": dir_loss_reduced, # 7:"cls_loss_reduced": cls_loss_reduced, # 8:"loc_loss_reduced": loc_loss_reduced, # 9:"cared": cared, # } # cls_preds = ret_dict["cls_preds"] cls_preds = ret_dict[5] # loss = ret_dict["loss"].mean() loss = ret_dict[0].mean() # cls_loss_reduced = ret_dict["cls_loss_reduced"].mean() cls_loss_reduced = ret_dict[7].mean() # loc_loss_reduced = ret_dict["loc_loss_reduced"].mean() loc_loss_reduced = ret_dict[8].mean() # cls_pos_loss = ret_dict["cls_pos_loss"] cls_pos_loss = ret_dict[3] # cls_neg_loss = ret_dict["cls_neg_loss"] cls_neg_loss = ret_dict[4] # loc_loss = ret_dict["loc_loss"] loc_loss = ret_dict[2] # cls_loss = ret_dict["cls_loss"] cls_loss = ret_dict[1] # dir_loss_reduced = ret_dict["dir_loss_reduced"] dir_loss_reduced = ret_dict[6] # cared = ret_dict["cared"] cared = ret_dict[9] # labels = example_torch["labels"] labels = example_tuple[8] if train_cfg.enable_mixed_precision: loss *= loss_scale loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) mixed_optimizer.step() mixed_optimizer.zero_grad() net.update_global_step() net_metrics = net.update_metrics(cls_loss_reduced, loc_loss_reduced, cls_preds, labels, cared) step_time = (time.time() - t) t = time.time() metrics = {} num_pos = int((labels > 0)[0].float().sum().cpu().numpy()) num_neg = int((labels == 0)[0].float().sum().cpu().numpy()) # if 'anchors_mask' not in example_torch: # num_anchors = example_torch['anchors'].shape[1] # else: # num_anchors = int(example_torch['anchors_mask'][0].sum()) num_anchors = int(example_tuple[7][0].sum()) global_step = net.get_global_step() if global_step % display_step == 0: loc_loss_elem = [ float(loc_loss[:, :, i].sum().detach().cpu().numpy() / batch_size) for i in range(loc_loss.shape[-1]) ] metrics["step"] = global_step metrics["steptime"] = step_time metrics.update(net_metrics) metrics["loss"] = {} metrics["loss"]["loc_elem"] = loc_loss_elem metrics["loss"]["cls_pos_rt"] = float( cls_pos_loss.detach().cpu().numpy()) metrics["loss"]["cls_neg_rt"] = float( cls_neg_loss.detach().cpu().numpy()) # if unlabeled_training: # metrics["loss"]["diff_rt"] = float( # diff_loc_loss_reduced.detach().cpu().numpy()) if model_cfg.use_direction_classifier: metrics["loss"]["dir_rt"] = float( dir_loss_reduced.detach().cpu().numpy()) # metrics["num_vox"] = int(example_torch["voxels"].shape[0]) metrics["num_vox"] = int(example_tuple[0].shape[0]) metrics["num_pos"] = int(num_pos) metrics["num_neg"] = int(num_neg) metrics["num_anchors"] = int(num_anchors) metrics["lr"] = float( mixed_optimizer.param_groups[0]['lr']) # metrics["image_idx"] = example['image_idx'][0] metrics["image_idx"] = example_tuple[11][0] flatted_metrics = flat_nested_json_dict(metrics) flatted_summarys = flat_nested_json_dict(metrics, "/") for k, v in flatted_summarys.items(): if isinstance(v, (list, tuple)): v = {str(i): e for i, e in enumerate(v)} writer.add_scalars(k, v, global_step) else: writer.add_scalar(k, v, global_step) metrics_str_list = [] for k, v in flatted_metrics.items(): if isinstance(v, float): metrics_str_list.append(f"{k}={v:.3}") elif isinstance(v, (list, tuple)): if v and isinstance(v[0], float): v_str = ', '.join([f"{e:.3}" for e in v]) metrics_str_list.append(f"{k}=[{v_str}]") else: metrics_str_list.append(f"{k}={v}") else: metrics_str_list.append(f"{k}={v}") log_str = ', '.join(metrics_str_list) print(log_str, file=logf) print(log_str) ckpt_elasped_time = time.time() - ckpt_start_time if ckpt_elasped_time > train_cfg.save_checkpoints_secs: torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) ckpt_start_time = time.time() total_step_elapsed += steps torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) # Ensure that all evaluation points are saved forever torchplus.train.save_models(eval_checkpoint_dir, [net, optimizer], net.get_global_step(), max_to_keep=100) # net.eval() # result_path_step = result_path / f"step_{net.get_global_step()}" # result_path_step.mkdir(parents=True, exist_ok=True) # print("#################################") # print("#################################", file=logf) # print("# EVAL") # print("# EVAL", file=logf) # print("#################################") # print("#################################", file=logf) # print("Generate output labels...") # print("Generate output labels...", file=logf) # t = time.time() # dt_annos = [] # prog_bar = ProgressBar() # prog_bar.start(len(eval_dataset) // eval_input_cfg.batch_size + 1) # for example in iter(eval_dataloader): # example = example_convert_to_torch(example, float_dtype) # # evaluation example:[0:'voxels', 1:'num_points', 2:'coordinates', 3:'rect', # # 4:'Trv2c', 5:'P2', # # 6:'anchors', 7:'anchors_mask', 8:'image_idx', 9:'image_shape'] # example_tuple = list(example.values()) # example_tuple[8] = torch.from_numpy(example_tuple[8]) # example_tuple[9] = torch.from_numpy(example_tuple[9]) # if pickle_result: # dt_annos += predict_kitti_to_anno( # net, example_tuple, class_names, center_limit_range, # model_cfg.lidar_input) # else: # _predict_kitti_to_file(net, example, result_path_step, # class_names, center_limit_range, # model_cfg.lidar_input) # # prog_bar.print_bar() # # sec_per_ex = len(eval_dataset) / (time.time() - t) # print(f"avg forward time per example: {net.avg_forward_time:.3f}") # print( # f"avg postprocess time per example: {net.avg_postprocess_time:.3f}" # ) # # net.clear_time_metrics() # print(f'generate label finished({sec_per_ex:.2f}/s). start eval:') # print( # f'generate label finished({sec_per_ex:.2f}/s). start eval:', # file=logf) # gt_annos = [ # info["annos"] for info in eval_dataset.dataset.kitti_infos # ] # if not pickle_result: # dt_annos = kitti.get_label_annos(result_path_step) # result, mAPbbox, mAPbev, mAP3d, mAPaos = get_official_eval_result(gt_annos, dt_annos, class_names, # return_data=True) # print(result, file=logf) # print(result) # writer.add_text('eval_result', result, global_step) # # for i, class_name in enumerate(class_names): # writer.add_scalar('bev_ap:{}'.format(class_name), mAPbev[i, 1, 0], global_step) # writer.add_scalar('3d_ap:{}'.format(class_name), mAP3d[i, 1, 0], global_step) # writer.add_scalar('aos_ap:{}'.format(class_name), mAPaos[i, 1, 0], global_step) # writer.add_scalar('bev_map', np.mean(mAPbev[:, 1, 0]), global_step) # writer.add_scalar('3d_map', np.mean(mAP3d[:, 1, 0]), global_step) # writer.add_scalar('aos_map', np.mean(mAPaos[:, 1, 0]), global_step) # # result = get_coco_eval_result(gt_annos, dt_annos, class_names) # print(result, file=logf) # print(result) # if pickle_result: # with open(result_path_step / "result.pkl", 'wb') as f: # pickle.dump(dt_annos, f) # writer.add_text('eval_result', result, global_step) # net.train() except Exception as e: torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) logf.close() raise e # save model before exit torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) logf.close()
def train(): """train a VoxelNet model specified by a config file. """ global args args = parse() create_folder = args.create_folder config_path = args.config_path model_dir = args.model_dir use_fusion = args.use_fusion use_ft = args.use_ft use_second_stage = args.use_second_stage result_path = args.result_path display_step = args.display_step summary_step = args.summary_step local_rank = args.local_rank pickle_result = args.pickle_result patchs = args.patchs # import pdb; pdb.set_trace() if create_folder: if pathlib.Path(model_dir).exists(): model_dir = torchplus.train.create_folder(model_dir) patchs = patchs or [] model_dir = pathlib.Path(model_dir) model_dir.mkdir(parents=True, exist_ok=True) if result_path is None: result_path = model_dir / 'results' config_file_bkp = "pipeline.config" config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) for patch in patchs: patch = "config." + patch exec(patch) shutil.copyfile(config_path, str(model_dir / config_file_bkp)) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config ################ Apex multi-gpu setting ################ if mp.get_start_method(allow_none=True) is None: mp.set_start_method('spawn') num_gpus = torch.cuda.device_count() torch.cuda.set_device(args.local_rank % num_gpus) dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:%d' % args.m_port, rank=args.local_rank, world_size=num_gpus) # assert batch_size % num_gpus == 0, 'Batch size should be matched with GPUS: (%d, %d)' % (batch_size, num_gpus) # batch_size_each_gpu = batch_size // num_gpus rank = dist.get_rank() # args_gpu = args.local_rank # torch.cuda.set_device(args_gpu) # torch.distributed.init_process_group(backend='nccl', init_method='env://') ######################################################## ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) ###################### # BUILD TARGET ASSIGNER ###################### bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) class_names = target_assigner.classes ###################### # BUILD NET ###################### center_limit_range = model_cfg.post_center_limit_range if use_second_stage: net = second_2stage_builder.build(model_cfg, voxel_generator, target_assigner) else: net = second_builder.build(model_cfg, voxel_generator, target_assigner) net.cuda() net.train() net = nn.parallel.DistributedDataParallel( net, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True) # net = DDP(net, delay_allreduce=True) if args.local_rank == 0: print("num_trainable parameters:", len(list(net.parameters()))) # for n, p in net.named_parameters(): # print(n, p.shape) ###################### if use_second_stage: if use_fusion: pth_name = 'pre_weight/8020/voxelnet-35210.tckpt' # pth_name= 'logs/test/voxelnet-35210.tckpt' if args.local_rank == 0: for i in range(30): print( '################## load Fusion First stage weight complete #######################' ) else: pth_name = 'pre_weight/first_stage/lidaronly/voxelnet-30950.tckpt' if args.local_rank == 0: for i in range(30): print( '################## load LiDAR Only First stage weight complete #######################' ) res_pre_weights = torch.load(pth_name) new_res_state_dict = OrderedDict() model_dict = net.state_dict() for k, v in res_pre_weights.items(): if 'global_step' not in k: if 'dir' not in k: k = 'module.' + k new_res_state_dict[k] = v # import pdb; pdb.set_trace() model_dict.update(new_res_state_dict) net.load_state_dict(model_dict) if use_fusion: # use_fusion: fpn_depth = 18 pth_name = 'pre_weight/FPN' + str(fpn_depth) + '_retinanet_968.pth' res_pre_weights = torch.load(pth_name) new_res_state_dict = OrderedDict() model_dict = net.state_dict() for k, v in res_pre_weights['state_dict'].items(): if ('regressionModel' not in k) and ('classificationModel' not in k): name = k.replace('module', 'rpn') new_res_state_dict[name] = v model_dict.update(new_res_state_dict) net.load_state_dict(model_dict) for i in range(30): print('!!!!!!!!!!!!!!!!!! load FPN' + str(fpn_depth) + ' weight complete !!!!!!!!!!!!!!!!!!') ############ load FPN18 pre-weight ############# if use_fusion and not use_second_stage: fpn_depth = 18 pth_name = 'pre_weight/FPN' + str(fpn_depth) + '_retinanet_968.pth' res_pre_weights = torch.load(pth_name) new_res_state_dict = OrderedDict() model_dict = net.module.state_dict() for k, v in res_pre_weights['state_dict'].items(): if ('regressionModel' not in k) and ('classificationModel' not in k): name = k.replace('module', 'rpn') new_res_state_dict[name] = v model_dict.update(new_res_state_dict) net.module.load_state_dict(model_dict) if args.local_rank == 0: for i in range(30): print('!!!!!!!!!!!!!!!!!! load FPN' + str(fpn_depth) + ' weight complete !!!!!!!!!!!!!!!!!!') ################################################ # BUILD OPTIMIZER ###################### # we need global_step to create lr_scheduler, so restore net first. # import pdb; pdb.set_trace() torchplus.train.try_restore_latest_checkpoints(model_dir, [net.module]) gstep = net.module.get_global_step() - 1 optimizer_cfg = train_cfg.optimizer if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) loss_scale = train_cfg.loss_scale_factor mixed_optimizer = optimizer_builder.build( optimizer_cfg, net, mixed=train_cfg.enable_mixed_precision, loss_scale=loss_scale) optimizer = mixed_optimizer """ if train_cfg.enable_mixed_precision: mixed_optimizer = torchplus.train.MixedPrecisionWrapper( optimizer, loss_scale) else: mixed_optimizer = optimizer """ # must restore optimizer AFTER using MixedPrecisionWrapper torchplus.train.try_restore_latest_checkpoints(model_dir, [mixed_optimizer]) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, optimizer, train_cfg.steps) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 ###################### # PREPARE INPUT ###################### dataset = input_reader_builder.build(input_cfg, model_cfg, training=True, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataset = input_reader_builder.build(eval_input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) def _worker_init_fn(worker_id): time_seed = np.array(time.time(), dtype=np.int32) np.random.seed(time_seed + worker_id) print(f"WORKER {worker_id} seed:", np.random.get_state()[1][0]) sampler = torch.utils.data.distributed.DistributedSampler(dataset) dataloader = torch.utils.data.DataLoader(dataset, batch_size=input_cfg.batch_size, shuffle=False, num_workers=input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch, worker_init_fn=_worker_init_fn, sampler=sampler) eval_sampler = torch.utils.data.distributed.DistributedSampler( eval_dataset) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=eval_input_cfg.batch_size, shuffle=False, num_workers=eval_input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch, sampler=eval_sampler) data_iter = iter(dataloader) print(len(dataloader)) # import pdb; pdb.set_trace() ###################### # TRAINING ###################### training_detail = [] log_path = model_dir / 'log.txt' training_detail_path = model_dir / 'log.json' if training_detail_path.exists(): with open(training_detail_path, 'r') as f: training_detail = json.load(f) logf = open(log_path, 'a') logf.write(proto_str) logf.write("\n") summary_dir = model_dir / 'summary' summary_dir.mkdir(parents=True, exist_ok=True) if args.local_rank == 0: writer = SummaryWriter(str(summary_dir)) else: writer = None total_step_elapsed = 0 remain_steps = train_cfg.steps - net.module.get_global_step() t = time.time() ckpt_start_time = t total_loop = train_cfg.steps // train_cfg.steps_per_eval + 1 # total_loop = remain_steps // train_cfg.steps_per_eval + 1 clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch if train_cfg.steps % train_cfg.steps_per_eval == 0: total_loop -= 1 mixed_optimizer.zero_grad() try: for _ in range(total_loop): if total_step_elapsed + train_cfg.steps_per_eval > train_cfg.steps: steps = train_cfg.steps % train_cfg.steps_per_eval else: steps = train_cfg.steps_per_eval for step in range(steps): lr_scheduler.step(net.module.get_global_step()) try: example = next(data_iter) except StopIteration: print("end epoch") if clear_metrics_every_epoch: net.module.clear_metrics() data_iter = iter(dataloader) example = next(data_iter) example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] ret_dict = net(example_torch) # box_preds = ret_dict["box_preds"] cls_preds = ret_dict["cls_preds"] loss = ret_dict["loss"].mean() cls_loss_reduced = ret_dict["cls_loss_reduced"].mean() loc_loss_reduced = ret_dict["loc_loss_reduced"].mean() cls_pos_loss = ret_dict["cls_pos_loss"] cls_neg_loss = ret_dict["cls_neg_loss"] loc_loss = ret_dict["loc_loss"] cls_loss = ret_dict["cls_loss"] # dir_loss_reduced = ret_dict["dir_loss_reduced"] cared = ret_dict["cared"] # labels = example_torch["labels"] if use_second_stage: labels = ret_dict["labels"] else: labels = example_torch["labels"] if train_cfg.enable_mixed_precision: loss *= loss_scale loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) mixed_optimizer.step() mixed_optimizer.zero_grad() net.module.update_global_step() net_metrics = net.module.update_metrics( cls_loss_reduced, loc_loss_reduced, cls_preds, labels, cared) step_time = (time.time() - t) t = time.time() metrics = {} num_pos = int((labels > 0)[0].float().sum().cpu().numpy()) num_neg = int((labels == 0)[0].float().sum().cpu().numpy()) if 'anchors_mask' not in example_torch: num_anchors = example_torch['anchors'].shape[1] else: num_anchors = int(example_torch['anchors_mask'][0].sum()) global_step = net.module.get_global_step() # if global_step % display_step == 0: if (writer is not None) and (global_step % display_step == 0): loc_loss_elem = [ float(loc_loss[:, :, i].sum().detach().cpu().numpy() / batch_size) for i in range(loc_loss.shape[-1]) ] metrics["type"] = "step_info" metrics["step"] = global_step metrics["steptime"] = step_time metrics.update(net_metrics) metrics["loss"] = {} metrics["loss"]["loc_elem"] = loc_loss_elem metrics["loss"]["cls_pos_rt"] = float( cls_pos_loss.detach().cpu().numpy()) metrics["loss"]["cls_neg_rt"] = float( cls_neg_loss.detach().cpu().numpy()) if model_cfg.use_direction_classifier: metrics["loss"]["dir_rt"] = float( dir_loss_reduced.detach().cpu().numpy()) metrics["num_vox"] = int(example_torch["voxels"].shape[0]) metrics["num_pos"] = int(num_pos) metrics["num_neg"] = int(num_neg) metrics["num_anchors"] = int(num_anchors) # metrics["lr"] = float( # mixed_optimizer.param_groups[0]['lr']) metrics["lr"] = float(optimizer.lr) metrics["image_idx"] = example['image_idx'][0] training_detail.append(metrics) flatted_metrics = flat_nested_json_dict(metrics) flatted_summarys = flat_nested_json_dict(metrics, "/") for k, v in flatted_summarys.items(): if isinstance(v, (list, tuple)): v = {str(i): e for i, e in enumerate(v)} if type(v) != str and ('loc_elem' not in k): writer.add_scalars(k, v, global_step) else: if (type(v) != str) and ('loc_elem' not in k): writer.add_scalar(k, v, global_step) metrics_str_list = [] for k, v in flatted_metrics.items(): if isinstance(v, float): metrics_str_list.append(f"{k}={v:.3}") elif isinstance(v, (list, tuple)): if v and isinstance(v[0], float): v_str = ', '.join([f"{e:.3}" for e in v]) metrics_str_list.append(f"{k}=[{v_str}]") else: metrics_str_list.append(f"{k}={v}") else: metrics_str_list.append(f"{k}={v}") log_str = ', '.join(metrics_str_list) print(log_str, file=logf) print(log_str) ckpt_elasped_time = time.time() - ckpt_start_time if ckpt_elasped_time > train_cfg.save_checkpoints_secs: if writer is not None: torchplus.train.save_models( model_dir, [net.module, optimizer], net.module.get_global_step()) ckpt_start_time = time.time() total_step_elapsed += steps if writer is not None: to_cpu = True if isinstance( net, torch.nn.parallel.DistributedDataParallel) else False # import pdb; pdb.set_trace() torchplus.train.save_models(model_dir, [net.module, optimizer], net.module.get_global_step(), to_cpu=to_cpu) net.eval() with torch.no_grad(): result_path_step = result_path / f"step_{net.module.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) print("#################################") print("#################################", file=logf) print("# EVAL") print("# EVAL", file=logf) print("#################################") print("#################################", file=logf) print("Generate output labels...") print("Generate output labels...", file=logf) t = time.time() dt_annos = [] prog_bar = ProgressBar() net.module.clear_timer() prog_bar.start( (len(eval_dataset) + eval_input_cfg.batch_size - 1) // eval_input_cfg.batch_size) for example in iter(eval_dataloader): example = example_convert_to_torch( example, float_dtype) if pickle_result: dt_annos += predict_kitti_to_anno( net.module, example, class_names, center_limit_range, model_cfg.lidar_input) else: _predict_kitti_to_file(net.moudule, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) print( f'generate label finished({sec_per_ex:.2f}/s). start eval:' ) print( f'generate label finished({sec_per_ex:.2f}/s). start eval:', file=logf) gt_annos = [ info["annos"] for info in eval_dataset.dataset.kitti_infos ] if not pickle_result: dt_annos = kitti.get_label_annos(result_path_step) # result = get_official_eval_result_v2(gt_annos, dt_annos, class_names) # print(json.dumps(result, indent=2), file=logf) # if writer is not None: result = get_official_eval_result(gt_annos, dt_annos, class_names) print(result, file=logf) print(result) result_1 = result.split("\n")[:5] result_2 = result.split("\n")[10:15] result_3 = result.split("\n")[20:25] emh = ['0_easy', '1_mod', '2_hard'] result_save = result_1 for i in range(len(result_save) - 1): save_targ = result_save[i + 1] name_val = save_targ.split(':')[0].split(' ')[0] value_val = save_targ.split(':')[1:] for ev in range(3): each_val = value_val[0].split(',')[ev] merge_txt = 'AP_kitti/car_70/' + name_val + '/' + emh[ ev] writer.add_scalar(merge_txt, float(each_val), global_step) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f) writer.add_text('eval_result', result, global_step) net.train() except Exception as e: if writer is not None: to_cpu = True if isinstance( net, torch.nn.parallel.DistributedDataParallel) else False torchplus.train.save_models(model_dir, [net, optimizer], net.module.get_global_step(), to_cpu=to_cpu) logf.close() raise e # save model before exit if writer is not None: to_cpu = True if isinstance( net, torch.nn.parallel.DistributedDataParallel) else False torchplus.train.save_models(model_dir, [net, optimizer], net.module.get_global_step(), to_cpu=to_cpu) logf.close()
def train(config_path, model_dir, result_path=None, create_folder=False, display_step=50, summary_step=5, pickle_result=True): """train a VoxelNet model specified by a config file. """ if create_folder: if pathlib.Path(model_dir).exists(): model_dir = torchplus.train.create_folder(model_dir) model_dir = pathlib.Path(model_dir) model_dir.mkdir(parents=True, exist_ok=True) if result_path is None: result_path = model_dir / 'results' config_file_bkp = "pipeline.config" config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) shutil.copyfile(config_path, str(model_dir / config_file_bkp)) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config class_names = list(input_cfg.class_names) ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) ###################### # BUILD TARGET ASSIGNER ###################### bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) ###################### # BUILD NET ###################### center_limit_range = model_cfg.post_center_limit_range net = second_builder_new.build(model_cfg, voxel_generator, target_assigner) net.cuda() # net_train = torch.nn.DataParallel(net).cuda() print("num_trainable parameters:", len(list(net.parameters()))) # for n, p in net.named_parameters(): # print(n, p.shape) ###################### # BUILD OPTIMIZER ###################### # we need global_step to create lr_scheduler, so restore net first. torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) gstep = net.get_global_step() - 1 optimizer_cfg = train_cfg.optimizer if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) optimizer = optimizer_builder.build(optimizer_cfg, net.parameters()) if train_cfg.enable_mixed_precision: loss_scale = train_cfg.loss_scale_factor mixed_optimizer = torchplus.train.MixedPrecisionWrapper( optimizer, loss_scale) else: mixed_optimizer = optimizer # must restore optimizer AFTER using MixedPrecisionWrapper torchplus.train.try_restore_latest_checkpoints(model_dir, [mixed_optimizer]) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, optimizer, gstep) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 ###################### # PREPARE INPUT ###################### dataset = input_reader_builder.build(input_cfg, model_cfg, training=True, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataset = input_reader_builder.build(eval_input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) def _worker_init_fn(worker_id): time_seed = np.array(time.time(), dtype=np.int32) np.random.seed(time_seed + worker_id) print(f"WORKER {worker_id} seed:", np.random.get_state()[1][0]) dataloader = torch.utils.data.DataLoader(dataset, batch_size=input_cfg.batch_size, shuffle=True, num_workers=input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch, worker_init_fn=_worker_init_fn) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=eval_input_cfg.batch_size, shuffle=False, num_workers=eval_input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch) data_iter = iter(dataloader) ###################### # TRAINING ###################### log_path = model_dir / 'log.txt' logf = open(log_path, 'a') logf.write(proto_str) logf.write("\n") summary_dir = model_dir / 'summary' summary_dir.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(str(summary_dir)) total_step_elapsed = 0 remain_steps = train_cfg.steps - net.get_global_step() t = time.time() ckpt_start_time = t total_loop = train_cfg.steps // train_cfg.steps_per_eval + 1 # total_loop = remain_steps // train_cfg.steps_per_eval + 1 clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch if train_cfg.steps % train_cfg.steps_per_eval == 0: total_loop -= 1 mixed_optimizer.zero_grad() try: for loop in range(total_loop): if total_step_elapsed + train_cfg.steps_per_eval > train_cfg.steps: steps = train_cfg.steps % train_cfg.steps_per_eval else: steps = train_cfg.steps_per_eval for step in range(steps): if check_if_should_pause(): torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) print('pause and save model @ {}/{} steps:{}'.format( loop, total_loop, net.get_global_step())) sys.exit(0) lr_scheduler.step() try: example = next(data_iter) except StopIteration: print("end epoch") if clear_metrics_every_epoch: net.clear_metrics() data_iter = iter(dataloader) example = next(data_iter) example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] ret_dict = net(example_torch) # box_preds = ret_dict["box_preds"] cls_preds = ret_dict["cls_preds"] loss = ret_dict["loss"].mean() cls_loss_reduced = ret_dict["cls_loss_reduced"].mean() loc_loss_reduced = ret_dict["loc_loss_reduced"].mean() cls_pos_loss = ret_dict["cls_pos_loss"] cls_neg_loss = ret_dict["cls_neg_loss"] loc_loss = ret_dict["loc_loss"] cls_loss = ret_dict["cls_loss"] dir_loss_reduced = ret_dict["dir_loss_reduced"] cared = ret_dict["cared"] labels = example_torch["labels"] if train_cfg.enable_mixed_precision: loss *= loss_scale loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) mixed_optimizer.step() mixed_optimizer.zero_grad() net.update_global_step() net_metrics = net.update_metrics(cls_loss_reduced, loc_loss_reduced, cls_preds, labels, cared) step_time = (time.time() - t) t = time.time() metrics = {} num_pos = int((labels > 0)[0].float().sum().cpu().numpy()) num_neg = int((labels == 0)[0].float().sum().cpu().numpy()) if 'anchors_mask' not in example_torch: num_anchors = example_torch['anchors'].shape[1] else: num_anchors = int(example_torch['anchors_mask'][0].sum()) global_step = net.get_global_step() if global_step % display_step == 0: loc_loss_elem = [ float(loc_loss[:, :, i].sum().detach().cpu().numpy() / batch_size) for i in range(loc_loss.shape[-1]) ] metrics["step"] = global_step metrics["steptime"] = step_time metrics.update(net_metrics) metrics["loss"] = {} metrics["loss"]["loc_elem"] = loc_loss_elem metrics["loss"]["cls_pos_rt"] = float( cls_pos_loss.detach().cpu().numpy()) metrics["loss"]["cls_neg_rt"] = float( cls_neg_loss.detach().cpu().numpy()) # if unlabeled_training: # metrics["loss"]["diff_rt"] = float( # diff_loc_loss_reduced.detach().cpu().numpy()) if model_cfg.use_direction_classifier: metrics["loss"]["dir_rt"] = float( dir_loss_reduced.detach().cpu().numpy()) metrics["num_vox"] = int(example_torch["voxels"].shape[0]) metrics["num_pos"] = int(num_pos) metrics["num_neg"] = int(num_neg) metrics["num_anchors"] = int(num_anchors) metrics["lr"] = float( mixed_optimizer.param_groups[0]['lr']) metrics["image_idx"] = example['image_idx'][0] flatted_metrics = flat_nested_json_dict(metrics) flatted_summarys = flat_nested_json_dict(metrics, "/") for k, v in flatted_summarys.items(): if isinstance(v, (list, tuple)): v = {str(i): e for i, e in enumerate(v)} writer.add_scalars(k, v, global_step) else: writer.add_scalar(k, v, global_step) metrics_str_list = [] for k, v in flatted_metrics.items(): if isinstance(v, float): metrics_str_list.append(f"{k}={v:.3}") elif isinstance(v, (list, tuple)): if v and isinstance(v[0], float): v_str = ', '.join([f"{e:.3}" for e in v]) metrics_str_list.append(f"{k}=[{v_str}]") else: metrics_str_list.append(f"{k}={v}") else: metrics_str_list.append(f"{k}={v}") log_str = ', '.join(metrics_str_list) print(log_str, file=logf) print(log_str) bev_map = example['bev_map'][0][-1:] writer.add_image('BEV', colorize(bev_map), global_step) ckpt_elasped_time = time.time() - ckpt_start_time if ckpt_elasped_time > train_cfg.save_checkpoints_secs: torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) ckpt_start_time = time.time() # After one epoch total_step_elapsed += steps torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) net.eval() result_path_step = result_path / f"step_{net.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) print("#################################") print("#################################", file=logf) print("# EVAL") print("# EVAL", file=logf) print("#################################") print("#################################", file=logf) print("Generate output labels...") print("Generate output labels...", file=logf) t = time.time() dt_annos = [] prog_bar = ProgressBar() prog_bar.start(len(eval_dataset) // eval_input_cfg.batch_size + 1) for example in iter(eval_dataloader): example = example_convert_to_torch(example, float_dtype) if pickle_result: dt_annos += predict_kitti_to_anno(net, example, class_names, center_limit_range, model_cfg.lidar_input) else: _predict_kitti_to_file(net, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) print(f"avg forward time per example: {net.avg_forward_time:.3f}") print( f"avg postprocess time per example: {net.avg_postprocess_time:.3f}" ) net.clear_time_metrics() print(f'generate label finished({sec_per_ex:.2f}/s). start eval:') print(f'generate label finished({sec_per_ex:.2f}/s). start eval:', file=logf) gt_annos = [ info["annos"] for info in eval_dataset.dataset.kitti_infos ] if not pickle_result: dt_annos = kitti.get_label_annos(result_path_step) result = get_official_eval_result(gt_annos, dt_annos, class_names) print(result, file=logf) print(result) writer.add_text('eval_result', result, global_step) result = get_coco_eval_result(gt_annos, dt_annos, class_names) print(result, file=logf) print(result) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f) writer.add_text('eval_result', result, global_step) net.train() except Exception as e: torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) logf.close() raise e # save model before exit torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) logf.close()
input_cfg = config.eval_input_reader #model_cfg: structure and loss parameters of net model_cfg = config.model.second #train_cfg: optimizer parameters and iteration steps train_cfg = config.train_config #class_names: ["Cyclist", "Pedestrian"] class_names = list(input_cfg.class_names) #post_center_limit_range: [0, -50, -2.5, 80, 50, -0.5] center_limit_range = model_cfg.post_center_limit_range ##generate voxel, initial anchors voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) # voxel_size = voxel_generator.voxel_size # pc_range = voxel_generator.point_cloud_range # grid_size = voxel_generator.grid_size # feature_map_size = grid_size[:2] // 2 # feature_map_size = [*feature_map_size, 1][::-1] # points = np.fromfile( # str(point_file), dtype=np.float32, # count=-1).reshape([-1, 4]) # voxels, coordinates, num_points = voxel_generator.generate( # points, 20000) # ret = target_assigner.generate_anchors(feature_map_size) # anchors = ret["anchors"]
def evaluate(config_path, model_dir, result_path=None, predict_test=False, ckpt_path=None, ref_detfile=None, pickle_result=True, measure_time=False, batch_size=None): model_dir = pathlib.Path(model_dir) if predict_test: result_name = 'predict_test' else: result_name = 'eval_results' if result_path is None: result_path = model_dir / result_name else: result_path = pathlib.Path(result_path) config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config center_limit_range = model_cfg.post_center_limit_range ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) class_names = target_assigner.classes net = second_builder.build(model_cfg, voxel_generator, target_assigner, measure_time=measure_time) net.cuda() if ckpt_path is None: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) else: torchplus.train.restore(ckpt_path, net) if train_cfg.enable_mixed_precision: net.half() print("half inference!") net.metrics_to_float() net.convert_norm_to_float(net) batch_size = batch_size or input_cfg.batch_size start = time.clock() eval_dataset = input_reader_builder.build(input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=batch_size, shuffle=False, num_workers=0, # input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch) elapsed = (time.clock() - start) print("Time used:", elapsed) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 net.eval() result_path_step = result_path / f"step_{net.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) t = time.time() dt_annos = [] global_set = None print("Generate output labels...") bar = ProgressBar() bar.start((len(eval_dataset) + batch_size - 1) // batch_size) prep_example_times = [] prep_times = [] t2 = time.time() # print("eval_dataloader.type", type(eval_dataloader)) # print("eval_dataloader.shape", len(eval_dataloader)) for example in iter(eval_dataloader): # print("example keys",example.keys()) # print("example type", type(example)) # print("example len", len(example)) # print("example.voxel.shape", example.get("voxels").shape) # print("example", example) if measure_time: prep_times.append(time.time() - t2) t1 = time.time() torch.cuda.synchronize() example = example_convert_to_torch(example, float_dtype) if measure_time: torch.cuda.synchronize() prep_example_times.append(time.time() - t1) if pickle_result: dt_annos += predict_kitti_to_anno(net, example, class_names, center_limit_range, model_cfg.lidar_input, global_set) else: _predict_kitti_to_file(net, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input) # print(json.dumps(net.middle_feature_extractor.middle_conv.sparity_dict)) bar.print_bar() if measure_time: t2 = time.time() sec_per_example = len(eval_dataset) / (time.time() - t) print(f'generate label finished({sec_per_example:.2f}/s). start eval:') if measure_time: print( f"avg example to torch time: {np.mean(prep_example_times) * 1000:.3f} ms" ) print(f"avg prep time: {np.mean(prep_times) * 1000:.3f} ms") for name, val in net.get_avg_time_dict().items(): print(f"avg {name} time = {val * 1000:.3f} ms") if not predict_test: gt_annos = [info["annos"] for info in eval_dataset.dataset.kitti_infos] if not pickle_result: dt_annos = kitti.get_label_annos(result_path_step) result = get_official_eval_result(gt_annos, dt_annos, class_names) # print(json.dumps(result, indent=2)) print(result) result = get_coco_eval_result(gt_annos, dt_annos, class_names) print(result) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f)
def test_onnx_for_trt(onnx_path, config_path, model_dir, ckpt_path=None): dummy_dev_pillar_x_ = np.random.random(size=(1, 1, 12000, 100)).astype(np.float32) dummy_dev_pillar_y_ = np.random.random(size=(1, 1, 12000, 100)).astype(np.float32) dummy_dev_pillar_z_ = np.random.random(size=(1, 1, 12000, 100)).astype(np.float32) dummy_dev_pillar_i_ = np.random.random(size=(1, 1, 12000, 100)).astype(np.float32) dummy_dev_num_points_per_pillar_ = np.random.random(size=(1, 1, 12000, 1)).astype( np.float32) dummy_dev_x_coors_for_sub_shaped_ = np.random.random(size=(1, 1, 12000, 100)).astype( np.float32) dummy_dev_y_coors_for_sub_shaped_ = np.random.random(size=(1, 1, 12000, 100)).astype( np.float32) dummy_dev_pillar_feature_mask_ = np.random.random(size=(1, 1, 12000, 100)).astype( np.float32) model = onnx.load(onnx_path) engine = backend.prepare(model, device='CUDA:0', max_batch_size=1) print("model read success") print() output_data = engine.run( (dummy_dev_pillar_x_, dummy_dev_pillar_y_, dummy_dev_pillar_z_, dummy_dev_pillar_i_, dummy_dev_num_points_per_pillar_, dummy_dev_x_coors_for_sub_shaped_, dummy_dev_y_coors_for_sub_shaped_, dummy_dev_pillar_feature_mask_)) # ##########compare with pytorch output ######################### for i in range(len(output_data)): print(output_data[i].shape) print(output_data[0][0, 0, 0:100]) model_dir = pathlib.Path(model_dir) config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) model_cfg = config.model.second voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) net = second_builder_for_official_onnx_and_cuda.build( model_cfg, voxel_generator, target_assigner) net.cuda() net.eval() # since the model is changed, dont restore first if ckpt_path is None: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) else: torchplus.train.restore(ckpt_path, net) dummy_dev_pillar_x_ = torch.as_tensor(dummy_dev_pillar_x_, device="cuda") dummy_dev_pillar_y_ = torch.as_tensor(dummy_dev_pillar_y_, device="cuda") dummy_dev_pillar_z_ = torch.as_tensor(dummy_dev_pillar_z_, device="cuda") dummy_dev_pillar_i_ = torch.as_tensor(dummy_dev_pillar_i_, device="cuda") dummy_dev_num_points_per_pillar_ = torch.as_tensor( dummy_dev_num_points_per_pillar_, device="cuda") dummy_dev_x_coors_for_sub_shaped_ = torch.as_tensor( dummy_dev_x_coors_for_sub_shaped_, device="cuda") dummy_dev_y_coors_for_sub_shaped_ = torch.as_tensor( dummy_dev_y_coors_for_sub_shaped_, device="cuda") dummy_dev_pillar_feature_mask_ = torch.as_tensor( dummy_dev_pillar_feature_mask_, device="cuda") output_pytorch = net.voxel_feature_extractor( dummy_dev_pillar_x_, dummy_dev_pillar_y_, dummy_dev_pillar_z_, dummy_dev_pillar_i_, dummy_dev_num_points_per_pillar_, dummy_dev_x_coors_for_sub_shaped_, dummy_dev_y_coors_for_sub_shaped_, dummy_dev_pillar_feature_mask_) print(output_pytorch[0, 0, 0:100])
def model_2_onnx(config_path, model_dir, ckpt_path=None): model_dir = pathlib.Path(model_dir) config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) model_cfg = config.model.second voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) net = second_builder_for_official_onnx_and_cuda.build( model_cfg, voxel_generator, target_assigner) # since the model is changed, dont restore first if ckpt_path is None: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) else: torchplus.train.restore(ckpt_path, net) # print(net) # convert model to onnx dummy_dev_pillar_x_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_pillar_y_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_pillar_z_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_pillar_i_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_num_points_per_pillar_ = torch.randn(1, 1, 12000, 1, device='cuda') dummy_dev_x_coors_for_sub_shaped_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_y_coors_for_sub_shaped_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_pillar_feature_mask_ = torch.randn(1, 1, 12000, 100, device='cuda') dummy_dev_scattered_feature = torch.randn(1, 64, 496, 432, device='cuda') net.cuda() net.eval() torch.onnx.export( net.voxel_feature_extractor, (dummy_dev_pillar_x_, dummy_dev_pillar_y_, dummy_dev_pillar_z_, dummy_dev_pillar_i_, dummy_dev_num_points_per_pillar_, dummy_dev_x_coors_for_sub_shaped_, dummy_dev_y_coors_for_sub_shaped_, dummy_dev_pillar_feature_mask_), "./pfe_test.onnx", verbose=False) torch.onnx.export(net.rpn, dummy_dev_scattered_feature, "./rpn_test.onnx", verbose=False)
def train(config_path, model_dir, result_path=None, create_folder=False, display_step=50, summary_step=5, pickle_result=True, patchs=None): torch.manual_seed(3) np.random.seed(3) if create_folder: if pathlib.Path(model_dir).exists(): model_dir = torchplus.train.create_folder(model_dir) patchs = patchs or [] model_dir = pathlib.Path(model_dir) model_dir.mkdir(parents=True, exist_ok=True) if result_path is None: result_path = model_dir / 'results' config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config detection_2d_path = config.train_config.detection_2d_path print("2d detection path:", detection_2d_path) center_limit_range = model_cfg.post_center_limit_range voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) class_names = target_assigner.classes net = build_inference_net('./configs/car.fhd.config', '../model_dir') fusion_layer = fusion.fusion() fusion_layer.cuda() optimizer_cfg = train_cfg.optimizer if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) loss_scale = train_cfg.loss_scale_factor mixed_optimizer = optimizer_builder.build( optimizer_cfg, fusion_layer, mixed=train_cfg.enable_mixed_precision, loss_scale=loss_scale) optimizer = mixed_optimizer # must restore optimizer AFTER using MixedPrecisionWrapper torchplus.train.try_restore_latest_checkpoints(model_dir, [mixed_optimizer]) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, optimizer, train_cfg.steps) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 ###################### # PREPARE INPUT ###################### dataset = input_reader_builder.build(input_cfg, model_cfg, training=True, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataset = input_reader_builder.build( eval_input_cfg, model_cfg, training=True, #if rhnning for test, here it needs to be False voxel_generator=voxel_generator, target_assigner=target_assigner) def _worker_init_fn(worker_id): time_seed = np.array(time.time(), dtype=np.int32) np.random.seed(time_seed + worker_id) print(f"WORKER {worker_id} seed:", np.random.get_state()[1][0]) dataloader = torch.utils.data.DataLoader(dataset, batch_size=input_cfg.batch_size, shuffle=True, num_workers=input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch, worker_init_fn=_worker_init_fn) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=eval_input_cfg.batch_size, shuffle=False, num_workers=eval_input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch) data_iter = iter(dataloader) ###################### # TRAINING ###################### focal_loss = SigmoidFocalClassificationLoss() cls_loss_sum = 0 training_detail = [] log_path = model_dir / 'log.txt' training_detail_path = model_dir / 'log.json' if training_detail_path.exists(): with open(training_detail_path, 'r') as f: training_detail = json.load(f) logf = open(log_path, 'a') logf.write(proto_str) logf.write("\n") summary_dir = model_dir / 'summary' summary_dir.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(str(summary_dir)) total_step_elapsed = 0 remain_steps = train_cfg.steps - net.get_global_step() t = time.time() ckpt_start_time = t total_loop = train_cfg.steps // train_cfg.steps_per_eval + 1 #print("steps, steps_per_eval, total_loop:", train_cfg.steps, train_cfg.steps_per_eval, total_loop) # total_loop = remain_steps // train_cfg.steps_per_eval + 1 clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch net.set_global_step(torch.tensor([0])) if train_cfg.steps % train_cfg.steps_per_eval == 0: total_loop -= 1 mixed_optimizer.zero_grad() try: for _ in range(total_loop): if total_step_elapsed + train_cfg.steps_per_eval > train_cfg.steps: steps = train_cfg.steps % train_cfg.steps_per_eval else: steps = train_cfg.steps_per_eval for step in range(steps): lr_scheduler.step(net.get_global_step()) try: example = next(data_iter) except StopIteration: print("end epoch") if clear_metrics_every_epoch: net.clear_metrics() data_iter = iter(dataloader) example = next(data_iter) example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] all_3d_output_camera_dict, all_3d_output, top_predictions, fusion_input, tensor_index = net( example_torch, detection_2d_path) d3_gt_boxes = example_torch["d3_gt_boxes"][0, :, :] if d3_gt_boxes.shape[0] == 0: target_for_fusion = np.zeros((1, 70400, 1)) positives = torch.zeros(1, 70400).type(torch.float32).cuda() negatives = torch.zeros(1, 70400).type(torch.float32).cuda() negatives[:, :] = 1 else: d3_gt_boxes_camera = box_torch_ops.box_lidar_to_camera( d3_gt_boxes, example_torch['rect'][0, :], example_torch['Trv2c'][0, :]) d3_gt_boxes_camera_bev = d3_gt_boxes_camera[:, [ 0, 2, 3, 5, 6 ]] ###### predicted bev boxes pred_3d_box = all_3d_output_camera_dict[0]["box3d_camera"] pred_bev_box = pred_3d_box[:, [0, 2, 3, 5, 6]] #iou_bev = bev_box_overlap(d3_gt_boxes_camera_bev.detach().cpu().numpy(), pred_bev_box.detach().cpu().numpy(), criterion=-1) iou_bev = d3_box_overlap( d3_gt_boxes_camera.detach().cpu().numpy(), pred_3d_box.squeeze().detach().cpu().numpy(), criterion=-1) iou_bev_max = np.amax(iou_bev, axis=0) #print(np.max(iou_bev_max)) target_for_fusion = ((iou_bev_max >= 0.7) * 1).reshape( 1, -1, 1) positive_index = ((iou_bev_max >= 0.7) * 1).reshape(1, -1) positives = torch.from_numpy(positive_index).type( torch.float32).cuda() negative_index = ((iou_bev_max <= 0.5) * 1).reshape(1, -1) negatives = torch.from_numpy(negative_index).type( torch.float32).cuda() cls_preds, flag = fusion_layer(fusion_input.cuda(), tensor_index.cuda()) one_hot_targets = torch.from_numpy(target_for_fusion).type( torch.float32).cuda() negative_cls_weights = negatives.type(torch.float32) * 1.0 cls_weights = negative_cls_weights + 1.0 * positives.type( torch.float32) pos_normalizer = positives.sum(1, keepdim=True).type( torch.float32) cls_weights /= torch.clamp(pos_normalizer, min=1.0) if flag == 1: cls_losses = focal_loss._compute_loss( cls_preds, one_hot_targets, cls_weights.cuda()) # [N, M] cls_losses_reduced = cls_losses.sum( ) / example_torch['labels'].shape[0] cls_loss_sum = cls_loss_sum + cls_losses_reduced if train_cfg.enable_mixed_precision: loss *= loss_scale cls_losses_reduced.backward() mixed_optimizer.step() mixed_optimizer.zero_grad() net.update_global_step() step_time = (time.time() - t) t = time.time() metrics = {} global_step = net.get_global_step() if global_step % display_step == 0: print("now it is", global_step, "steps", " and the cls_loss is :", cls_loss_sum / display_step, "learning_rate: ", float(optimizer.lr), file=logf) print("now it is", global_step, "steps", " and the cls_loss is :", cls_loss_sum / display_step, "learning_rate: ", float(optimizer.lr)) cls_loss_sum = 0 ckpt_elasped_time = time.time() - ckpt_start_time if ckpt_elasped_time > train_cfg.save_checkpoints_secs: torchplus.train.save_models(model_dir, [fusion_layer, optimizer], net.get_global_step()) ckpt_start_time = time.time() total_step_elapsed += steps torchplus.train.save_models(model_dir, [fusion_layer, optimizer], net.get_global_step()) fusion_layer.eval() net.eval() result_path_step = result_path / f"step_{net.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) print("#################################") print("#################################", file=logf) print("# EVAL") print("# EVAL", file=logf) print("#################################") print("#################################", file=logf) print("Generate output labels...") print("Generate output labels...", file=logf) t = time.time() dt_annos = [] prog_bar = ProgressBar() net.clear_timer() prog_bar.start( (len(eval_dataset) + eval_input_cfg.batch_size - 1) // eval_input_cfg.batch_size) val_loss_final = 0 for example in iter(eval_dataloader): example = example_convert_to_torch(example, float_dtype) if pickle_result: dt_annos_i, val_losses = predict_kitti_to_anno( net, detection_2d_path, fusion_layer, example, class_names, center_limit_range, model_cfg.lidar_input) dt_annos += dt_annos_i val_loss_final = val_loss_final + val_losses else: _predict_kitti_to_file(net, detection_2d_path, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) print("validation_loss:", val_loss_final / len(eval_dataloader)) print("validation_loss:", val_loss_final / len(eval_dataloader), file=logf) print(f'generate label finished({sec_per_ex:.2f}/s). start eval:') print(f'generate label finished({sec_per_ex:.2f}/s). start eval:', file=logf) gt_annos = [ info["annos"] for info in eval_dataset.dataset.kitti_infos ] if not pickle_result: dt_annos = kitti.get_label_annos(result_path_step) # result = get_official_eval_result_v2(gt_annos, dt_annos, class_names) result = get_official_eval_result(gt_annos, dt_annos, class_names) print(result, file=logf) print(result) writer.add_text('eval_result', json.dumps(result, indent=2), global_step) result = get_coco_eval_result(gt_annos, dt_annos, class_names) print(result, file=logf) print(result) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f) writer.add_text('eval_result', result, global_step) #net.train() fusion_layer.train() except Exception as e: torchplus.train.save_models(model_dir, [fusion_layer, optimizer], net.get_global_step()) logf.close() raise e # save model before exit torchplus.train.save_models(model_dir, [fusion_layer, optimizer], net.get_global_step()) logf.close()
def train(config_path, model_dir, result_path=None, create_folder=False, display_step=50, pretrained_path=None, pretrained_include=None, pretrained_exclude=None, freeze_include=None, freeze_exclude=None, multi_gpu=False, measure_time=False, resume=False): """train a PointPillars model specified by a config file. """ torch.cuda.empty_cache() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_dir = str(Path(model_dir).resolve()) if create_folder: if Path(model_dir).exists(): model_dir = torchplus.train.create_folder(model_dir) model_dir = Path(model_dir) if not resume and model_dir.exists(): raise ValueError("model dir exists and you don't specify resume.") model_dir.mkdir(parents=True, exist_ok=True) if result_path is None: result_path = model_dir / 'results' config, proto_str = load_config(model_dir, config_path) input_cfg = config.train_input_reader model_cfg = config.model.second train_cfg = config.train_config target_assigner_cfg = model_cfg.target_assigner voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) box_coder.custom_ndim = target_assigner._anchor_generators[0].custom_ndim net = PointPillarsNet(1, voxel_generator.grid_size, target_assigner.num_anchors_per_location, target_assigner.box_coder.code_size, with_distance=False).to(device) kaiming_init(net, 1.0) net_loss = build_net_loss(model_cfg, target_assigner).to(device) net_loss.clear_global_step() net_loss.clear_metrics() # print("num parameters:", len(list(net.parameters()))) load_pretrained_model(net, pretrained_path, pretrained_include, pretrained_exclude, freeze_include, freeze_exclude) if resume: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) amp_optimizer, lr_scheduler = create_optimizer(model_dir, train_cfg, net) collate_fn = merge_second_batch num_gpu = 1 ###################### # PREPARE INPUT ###################### dataset = input_reader_builder.build(input_cfg, model_cfg, training=True, voxel_generator=voxel_generator, target_assigner=target_assigner, multi_gpu=multi_gpu) dataloader = torch.utils.data.DataLoader( dataset, batch_size=input_cfg.batch_size * num_gpu, shuffle=True, num_workers=input_cfg.preprocess.num_workers * num_gpu, pin_memory=False, collate_fn=collate_fn, worker_init_fn=_worker_init_fn, drop_last=not multi_gpu) ###################### # TRAINING ###################### model_logging = SimpleModelLog(model_dir) model_logging.open() model_logging.log_text(proto_str + "\n", 0, tag="config") start_step = net_loss.get_global_step() total_step = train_cfg.steps t = time.time() steps_per_eval = train_cfg.steps_per_eval clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch amp_optimizer.zero_grad() step_times = [] step = start_step best_mAP = 0 epoch = 0 net.train() net_loss.train() try: while True: if clear_metrics_every_epoch: net_loss.clear_metrics() for example in dataloader: lr_scheduler.step(net_loss.get_global_step()) time_metrics = example["metrics"] example.pop("metrics") example_torch = example_convert_to_torch(example, float_dtype) batch_size = example_torch["anchors"].shape[0] coors = example_torch["coordinates"] input_features = compute_model_input( voxel_generator.voxel_size, voxel_generator.point_cloud_range, with_distance=False, voxels=example_torch['voxels'], num_voxels=example_torch['num_points'], coors=coors) # input_features = reshape_input(batch_size, input_features, coors, voxel_generator.grid_size) input_features = reshape_input1(input_features) net.batch_size = batch_size preds_list = net(input_features, coors) ret_dict = net_loss(example_torch, preds_list) cls_preds = ret_dict["cls_preds"] loss = ret_dict["loss"].mean() cls_loss_reduced = ret_dict["cls_loss_reduced"].mean() loc_loss_reduced = ret_dict["loc_loss_reduced"].mean() cls_pos_loss = ret_dict["cls_pos_loss"].mean() cls_neg_loss = ret_dict["cls_neg_loss"].mean() loc_loss = ret_dict["loc_loss"] cls_loss = ret_dict["cls_loss"] cared = ret_dict["cared"] labels = example_torch["labels"] loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) amp_optimizer.step() amp_optimizer.zero_grad() net_loss.update_global_step() net_metrics = net_loss.update_metrics(cls_loss_reduced, loc_loss_reduced, cls_preds, labels, cared) step_time = (time.time() - t) step_times.append(step_time) t = time.time() metrics = {} num_pos = int((labels > 0)[0].float().sum().cpu().numpy()) num_neg = int((labels == 0)[0].float().sum().cpu().numpy()) if 'anchors_mask' not in example_torch: num_anchors = example_torch['anchors'].shape[1] else: num_anchors = int(example_torch['anchors_mask'][0].sum()) global_step = net_loss.get_global_step() if global_step % display_step == 0: loc_loss_elem = [ float(loc_loss[:, :, i].sum().detach().cpu().numpy() / batch_size) for i in range(loc_loss.shape[-1]) ] metrics["runtime"] = { "step": global_step, "steptime": np.mean(step_times), } metrics["runtime"].update(time_metrics[0]) step_times = [] metrics.update(net_metrics) metrics["loss"]["loc_elem"] = loc_loss_elem metrics["loss"]["cls_pos_rt"] = float( cls_pos_loss.detach().cpu().numpy()) metrics["loss"]["cls_neg_rt"] = float( cls_neg_loss.detach().cpu().numpy()) if model_cfg.use_direction_classifier: dir_loss_reduced = ret_dict["dir_loss_reduced"].mean() metrics["loss"]["dir_rt"] = float( dir_loss_reduced.detach().cpu().numpy()) metrics["misc"] = { "num_vox": int(example_torch["voxels"].shape[0]), "num_pos": int(num_pos), "num_neg": int(num_neg), "num_anchors": int(num_anchors), "lr": float(amp_optimizer.lr), "mem_usage": psutil.virtual_memory().percent, } model_logging.log_metrics(metrics, global_step) step += 1 epoch += 1 if epoch % 2 == 0: global_step = net_loss.get_global_step() torchplus.train.save_models(model_dir, [net, amp_optimizer], global_step) net.eval() net_loss.eval() best_mAP = evaluate(net, net_loss, best_mAP, voxel_generator, target_assigner, config, model_logging, model_dir, result_path) net.train() net_loss.train() if epoch > 100: break if epoch > 100: break except Exception as e: print(json.dumps(example["metadata"], indent=2)) model_logging.log_text(str(e), step) model_logging.log_text(json.dumps(example["metadata"], indent=2), step) torchplus.train.save_models(model_dir, [net, amp_optimizer], step) raise e finally: model_logging.close() torchplus.train.save_models(model_dir, [net, amp_optimizer], net_loss.get_global_step())
def evaluate(config_path, model_dir, result_path=None, predict_test=False, ckpt_path=None, ref_detfile=None, pickle_result=True): model_dir = pathlib.Path(model_dir) if predict_test: result_name = 'predict_test' else: result_name = 'eval_results' if result_path is None: result_path = model_dir / result_name else: result_path = pathlib.Path(result_path) config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config class_names = list(input_cfg.class_names) center_limit_range = model_cfg.post_center_limit_range ###################### # BUILD VOXEL GENERATOR ###################### voxel_generator = voxel_builder.build(model_cfg.voxel_generator) bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] box_coder = box_coder_builder.build(model_cfg.box_coder) target_assigner_cfg = model_cfg.target_assigner target_assigner = target_assigner_builder.build(target_assigner_cfg, bv_range, box_coder) net = second_builder.build(model_cfg, voxel_generator, target_assigner) net.cuda() if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) if ckpt_path is None: torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) else: torchplus.train.restore(ckpt_path, net) eval_dataset = input_reader_builder.build(input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=input_cfg.batch_size, shuffle=False, num_workers=input_cfg.num_workers, pin_memory=False, collate_fn=merge_second_batch) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 net.eval() result_path_step = result_path / f"step_{net.get_global_step()}" result_path_step.mkdir(parents=True, exist_ok=True) t = time.time() if model_cfg.rpn.module_class_name == "PSA" or model_cfg.rpn.module_class_name == "RefineDet": dt_annos_coarse = [] dt_annos_refine = [] print("Generate output labels...") bar = ProgressBar() bar.start(len(eval_dataset) // input_cfg.batch_size + 1) for example in iter(eval_dataloader): example = example_convert_to_torch(example, float_dtype) if pickle_result: coarse, refine = predict_kitti_to_anno(net, example, class_names, center_limit_range, model_cfg.lidar_input, use_coarse_to_fine=True, global_set=None) dt_annos_coarse += coarse dt_annos_refine += refine else: _predict_kitti_to_file(net, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input, use_coarse_to_fine=True) bar.print_bar() else: dt_annos = [] print("Generate output labels...") bar = ProgressBar() bar.start(len(eval_dataset) // input_cfg.batch_size + 1) for example in iter(eval_dataloader): example = example_convert_to_torch(example, float_dtype) if pickle_result: dt_annos += predict_kitti_to_anno(net, example, class_names, center_limit_range, model_cfg.lidar_input, use_coarse_to_fine=False, global_set=None) else: _predict_kitti_to_file(net, example, result_path_step, class_names, center_limit_range, model_cfg.lidar_input, use_coarse_to_fine=False) bar.print_bar() sec_per_example = len(eval_dataset) / (time.time() - t) print(f'generate label finished({sec_per_example:.2f}/s). start eval:') print(f"avg forward time per example: {net.avg_forward_time:.3f}") print(f"avg postprocess time per example: {net.avg_postprocess_time:.3f}") if not predict_test: gt_annos = [info["annos"] for info in eval_dataset.dataset.kitti_infos] if not pickle_result: dt_annos = kitti.get_label_annos(result_path_step) if model_cfg.rpn.module_class_name == "PSA" or model_cfg.rpn.module_class_name == "RefineDet": print('Before Refine:') result_coarse = get_official_eval_result(gt_annos, dt_annos_coarse, class_names) print(result_coarse) print('After Refine:') result_refine = get_official_eval_result(gt_annos, dt_annos_refine, class_names) print(result_refine) result = get_coco_eval_result(gt_annos, dt_annos_refine, class_names) dt_annos = dt_annos_refine print(result) else: result = get_official_eval_result(gt_annos, dt_annos, class_names) print(result) result = get_coco_eval_result(gt_annos, dt_annos, class_names) print(result) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f)