def create_optimizer(model_dir, train_cfg, net): optimizer_cfg = train_cfg.optimizer loss_scale = train_cfg.loss_scale_factor fastai_optimizer = optimizer_builder.build(optimizer_cfg, net, mixed=False, loss_scale=loss_scale) amp_optimizer = fastai_optimizer torchplus.train.try_restore_latest_checkpoints(model_dir, [fastai_optimizer]) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, amp_optimizer, train_cfg.steps) return amp_optimizer, lr_scheduler
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.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] 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) 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) 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, 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(config_path, model_dir, result_path=None, create_folder=False, display_step=50, summary_step=5, 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 VoxelNet model specified by a config file. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # create dir for saving training states 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' # loadd config file config_file_bkp = "pipeline.config" if isinstance(config_path, str): # directly provide a config object. this usually used # when you want to train with several different parameters in # one script. config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) else: config = config_path proto_str = text_format.MessageToString(config, indent=2) with (model_dir / config_file_bkp).open("w") as f: f.write(proto_str) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config net = build_network(model_cfg, measure_time).to(device) # if train_cfg.enable_mixed_precision: # net.half() # net.metrics_to_float() # net.convert_norm_to_float(net) target_assigner = net.target_assigner voxel_generator = net.voxel_generator print("num parameters:", len(list(net.parameters()))) torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) if pretrained_path is not None: ## load pretrained params model_dict = net.state_dict() pretrained_dict = torch.load(pretrained_path) pretrained_dict = filter_param_dict(pretrained_dict, pretrained_include, pretrained_exclude) new_pretrained_dict = {} for k, v in pretrained_dict.items(): if k in model_dict and v.shape == model_dict[k].shape: new_pretrained_dict[k] = v print("Load pretrained parameters:") for k, v in new_pretrained_dict.items(): print(k, v.shape) model_dict.update(new_pretrained_dict) net.load_state_dict(model_dict) freeze_params_v2(dict(net.named_parameters()), freeze_include, freeze_exclude) net.clear_global_step() net.clear_metrics() if multi_gpu: net_parallel = torch.nn.DataParallel(net) else: net_parallel = net optimizer_cfg = train_cfg.optimizer loss_scale = train_cfg.loss_scale_factor fastai_optimizer = optimizer_builder.build(optimizer_cfg, net, mixed=False, loss_scale=loss_scale) if loss_scale < 0: loss_scale = "dynamic" if train_cfg.enable_mixed_precision: max_num_voxels = input_cfg.preprocess.max_number_of_voxels * input_cfg.batch_size assert max_num_voxels < 65535, "spconv fp16 training only support this" from apex import amp net, amp_optimizer = amp.initialize(net, fastai_optimizer, opt_level="O2", keep_batchnorm_fp32=True, loss_scale=loss_scale) net.metrics_to_float() else: amp_optimizer = fastai_optimizer torchplus.train.try_restore_latest_checkpoints(model_dir, [fastai_optimizer]) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, amp_optimizer, train_cfg.steps) if train_cfg.enable_mixed_precision: float_dtype = torch.float16 else: float_dtype = torch.float32 if multi_gpu: num_gpu = torch.cuda.device_count() print(f"MULTI-GPU: use {num_gpu} gpu") collate_fn = merge_second_batch_multigpu else: 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) eval_dataset = input_reader_builder.build(eval_input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) 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) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=eval_input_cfg.batch_size, # only support multi-gpu train shuffle=False, num_workers=eval_input_cfg.preprocess.num_workers, pin_memory=False, collate_fn=merge_second_batch) ###################### # TRAINING ###################### model_logging = SimpleModelLog(model_dir) model_logging.open() model_logging.log_text(proto_str + "\n", 0, tag="config") start_step = net.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 try: while True: if clear_metrics_every_epoch: net.clear_metrics() for example in dataloader: lr_scheduler.step(net.get_global_step()) time_metrics = example["metrics"] example.pop("metrics") example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] ret_dict = net_parallel(example_torch) 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"] if train_cfg.enable_mixed_precision: with amp.scale_loss(loss, amp_optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) amp_optimizer.step() amp_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) 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.get_global_step() if global_step % display_step == 0: if measure_time: for name, val in net.get_avg_time_dict().items(): print(f"avg {name} time = {val * 1000:.3f} ms") 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) if global_step % steps_per_eval == 0: torchplus.train.save_models(model_dir, [net, amp_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) model_logging.log_text("#################################", global_step) model_logging.log_text("# EVAL", global_step) model_logging.log_text("#################################", global_step) model_logging.log_text("Generate output labels...", global_step) t = time.time() detections = [] prog_bar = ProgressBar() net.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) detections += net(example) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) model_logging.log_text( f'generate label finished({sec_per_ex:.2f}/s). start eval:', global_step) result_dict = eval_dataset.dataset.evaluation( detections, str(result_path_step)) for k, v in result_dict["results"].items(): model_logging.log_text("Evaluation {}".format(k), global_step) model_logging.log_text(v, global_step) model_logging.log_metrics(result_dict["detail"], global_step) with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(detections, f) net.train() step += 1 if step >= total_step: break if step >= total_step: 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.get_global_step())
def train(config_path, model_dir, use_fusion=True, use_ft=False, use_second_stage=True, use_endtoend=True, result_path=None, create_folder=False, display_step=50, summary_step=5, local_rank=0, pickle_result=True, patchs=None): """train a VoxelNet mod[el specified by a config file. """ ############ tracking config_tr_path = '/mnt/new_iou/second.pytorch/second/mmMOT/experiments/second/spatio_test/config.yaml' load_tr_path = '/mnt/new_iou/second.pytorch/second/mmMOT/experiments/second/spatio_test/results' with open(config_tr_path) as f: config_tr = yaml.load(f, Loader=yaml.FullLoader) result_path_tr = load_tr_path config_tr = EasyDict(config_tr['common']) config_tr.save_path = os.path.dirname(config_tr_path) # create model # model_tr = build_model(config_tr) # model_tr.cuda() # optimizer_tr = build_optim(model_tr, config_tr) criterion_tr = build_criterion(config_tr.loss) last_iter = -1 best_mota = 0 # if load_tr_path: # if False: # best_mota, last_iter = load_state( # load_tr_path, model_tr, optimizer=optimizer_tr) # else: # load_state(load_tr_path, model_tr) cudnn.benchmark = True # Data loading code train_transform, valid_transform = build_augmentation(config_tr.augmentation) # # train # train_dataset = build_dataset( # config_tr, # set_source='train', # evaluate=False, # train_transform=train_transform) # trainval_dataset = build_dataset( # config_tr, # set_source='train', # evaluate=True, # valid_transform=valid_transform) # val_dataset = build_dataset( # config_tr, # set_source='val', # evaluate=True, # valid_transform=valid_transform) # train_sampler = DistributedGivenIterationSampler( # train_dataset, # config_tr.lr_scheduler.max_iter, # config_tr.batch_size, # world_size=1, # rank=0, # last_iter=last_iter) # import pdb; pdb.set_trace() # train_loader = DataLoader( # train_dataset, # batch_size=config_tr.batch_size, # shuffle=False, # num_workers=config_tr.workers, # pin_memory=True) tb_logger = SummaryWriter(config_tr.save_path + '/events') logger = create_logger('global_logger', config_tr.save_path + '/log.txt') # logger.info('args: {}'.format(pprint.pformat(args))) logger.info('config: {}'.format(pprint.pformat(config_tr))) # tracking_module = TrackingModule(model_tr, criterion_tr, # config_tr.det_type) # tracking_module.model.train() #### tracking setup done 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 ###################### # 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) if use_endtoend: net = second_endtoend_builder_spatio.build(model_cfg, voxel_generator, target_assigner, criterion_tr, config_tr.det_type) else: net = second_builder.build(model_cfg, voxel_generator, target_assigner) net.cuda() print("num_trainable parameters:", len(list(net.parameters()))) for n, p in net.named_parameters(): print(n, p.shape) # pth_name = './pre_weight/first_stage_gating_det/voxelnet-17013.tckpt' pth_name = './pre_weight/second_stage_gating_det/voxelnet-35000.tckpt' 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: new_res_state_dict[k] = v model_dict.update(new_res_state_dict) net.load_state_dict(model_dict) # for k, weight in dict(net.named_parameters()).items(): # lidar_conv, p_lidar_conv, fusion_module, w_det, w_link, appearance, point_net # if 'middle_feature_extractor' in '%s'%(k) or 'rpn' in '%s'%(k) or 'second_rpn' in '%s'%(k): # weight.requires_grad = False # 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) 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 # 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 ###################### # import pdb; pdb.set_trace() dataset = input_reader_builder_tr_vid_spatio.build( input_cfg, model_cfg, training=True, voxel_generator=voxel_generator, target_assigner=target_assigner, config_tr=config_tr, set_source='train', evaluate=False, train_transform=train_transform) eval_dataset = input_reader_builder_tr_vid_spatio.build( eval_input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner, config_tr=config_tr, set_source='val', evaluate=True, valid_transform=valid_transform) 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_tr_vid_spatio, 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_tr_vid_spatio) data_iter = iter(dataloader) ###################### # 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) 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 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() # optimizer_tr.zero_grad() logger = logging.getLogger('global_logger') best_mota = 0 losses = AverageMeter(config_tr.print_freq) total_steps = train_cfg.steps total_loop = total_steps // len(dataloader) kkkk = 0 for step in range(total_loop): for i, (example) in enumerate(dataloader): curr_step = 0 + i kkkk += 1 lr_scheduler.step(net.get_global_step()) example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] ret_dict = net(example_torch, train_param=True) 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"] # loss_tr = ret_dict["loss_tr"] if use_second_stage or use_endtoend: labels = ret_dict["labels"] else: labels = example_torch["labels"] if train_cfg.enable_mixed_precision: loss *= loss_scale try: loss.backward() except: abc = 1 # import pdb; pdb.set_trace() # abc = 1 # torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) # optimizer_tr.step() # optimizer_tr.zero_grad() 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() # print(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["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( optimizer.lr) metrics["image_idx"] = example['image_idx'][0][7:] 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: torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) ckpt_start_time = time.time() if kkkk > 0 and (kkkk) % config_tr.val_freq == 0: # if True: 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() net.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: results = predict_kitti_to_anno( net, example, class_names, center_limit_range, model_cfg.lidar_input) dt_annos += results 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'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) 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] try: writer.add_scalar(merge_txt, float(each_val), global_step) except: abc=1 import pdb; pdb.set_trace() abc=1 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) logger.info('Evaluation on validation set:') # MOTA, MOTP, recall, prec, F1, fp, fn, id_switches = validate( # val_dataset, # net, # str(0 + 1), # config_tr, # result_path_tr, # part='val') # print(MOTA, MOTP, recall, prec, F1, fp, fn, id_switches) # curr_step = step # if tb_logger is not None: # tb_logger.add_scalar('prec', prec, curr_step) # tb_logger.add_scalar('recall', recall, curr_step) # tb_logger.add_scalar('mota', MOTA, curr_step) # tb_logger.add_scalar('motp', MOTP, curr_step) # tb_logger.add_scalar('fp', fp, curr_step) # tb_logger.add_scalar('fn', fn, curr_step) # tb_logger.add_scalar('f1', F1, curr_step) # tb_logger.add_scalar('id_switches', id_switches, curr_step) # if lr_scheduler is not None: # tb_logger.add_scalar('lr', current_lr, curr_step) # is_best = MOTA > best_mota # best_mota = max(MOTA, best_mota) # print(best_mota) # import pdb; pdb.set_trace() # save_checkpoint( # { 'step': net.get_global_step(), # 'score_arch': config_tr.model.score_arch, # 'appear_arch': config_tr.model.appear_arch, # 'best_mota': best_mota, # 'state_dict': tracking_module.model.state_dict(), # 'optimizer': tracking_module.optimizer.state_dict(), # }, is_best, config_tr.save_path + '/ckpt') # net.train() # save model before exit torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) logf.close()
def train(config_path, model_dir, use_fusion=False, use_ft=False, use_second_stage=False, use_endtoend=False, result_path=None, create_folder=False, display_step=50, summary_step=5, local_rank=0, pickle_result=True, patchs=None): """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) 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 ###################### # 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) if use_endtoend: net = second_endtoend_builder.build(model_cfg, voxel_generator, target_assigner) else: net = second_builder.build(model_cfg, voxel_generator, target_assigner) net.cuda() # import pdb; pdb.set_trace() print("num_trainable parameters:", len(list(net.parameters()))) # for n, p in net.named_parameters(): # print(n, p.shape) # pth_name = 'pre_weight/first_stage/fusion_split/voxelnet-35210.tckpt' # # pth_name = 'pre_weight/first_stage/fusion_split/voxelnet-20130.tckpt' # 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: # new_res_state_dict[k] = v # model_dict.update(new_res_state_dict) # net.load_state_dict(model_dict) ###################### if use_second_stage or use_endtoend: if use_fusion: # pth_name = 'pre_weight/8020/voxelnet-20130.tckpt' pth_name = 'pre_weight/first_stage/fusion_split/voxelnet-35210.tckpt' for i in range(30): print( '################## load Fusion First stage weight complete #######################' ) else: pth_name = 'pre_weight/first_stage/lidaronly/voxelnet-30950.tckpt' 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: new_res_state_dict[k] = v model_dict.update(new_res_state_dict) net.load_state_dict(model_dict) ############ load FPN18 pre-weight ############# if (use_fusion and not use_second_stage and not use_endtoend): # if True: # or (use_endtoend and 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 !!!!!!!!!!!!!!!!!!') ################################################ # 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) 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]) 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 ###################### 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 # 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.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] 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"] # idx_offset = ret_dict["idx_offset"] # labels = example_torch["labels"] if use_second_stage or use_endtoend: labels = ret_dict["labels"] else: labels = example_torch["labels"] if train_cfg.enable_mixed_precision: loss *= loss_scale loss.backward() # import pdb; pdb.set_trace() 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() # print(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["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["idx_offset_mean"] = float(idx_offset.mean().detach().cpu().numpy()) # metrics["idx_offset_sum"] = float(idx_offset.sum().detach().cpu().numpy()) # 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) # if use_second_stage or use_endtoend: # bev_logs = ret_dict['bev_crops_output'][:64,0,...].view(64,1,14,14) # bev_vis = torchvision.utils.make_grid(bev_logs,normalize=True,scale_each=True) # writer.add_image('bev_crop',img_tensor=bev_vis, global_step=global_step) # if ret_dict['concat_crops_output'] is not None: # concat_logs = ret_dict['concat_crops_output'][:64,0,...].view(64,1,14,14) # concat_vis = torchvision.utils.make_grid(concat_logs,normalize=True,scale_each=True) # writer.add_image('concat_crop',img_tensor=concat_vis, global_step=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()) 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) 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'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) 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: 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(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 NetWork ###################### 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 write with training input size!" # ret_dict = net(example_torch) # Training Input form example 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, z from pillar center # assume config_file 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 ################################################################ # 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 write with training output size!" cls_preds = ret_dict[5] loss = ret_dict[0].mean() cls_loss_reduced = ret_dict[7].mean() loc_loss_reduced = ret_dict[8].mean() cls_pos_loss = ret_dict[3] cls_neg_loss = ret_dict[4] loc_loss = ret_dict[2] cls_loss = ret_dict[1] dir_loss_reduced = ret_dict[6] cared = ret_dict[9] 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_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_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) 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(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(config_path, model_dir, result_path=None, create_folder=False, display_step=50, summary_step=5, resume=False): """train a VoxelNet model specified by a config file. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if create_folder: if pathlib.Path(model_dir).exists(): model_dir = torchplus.train.create_folder(model_dir) model_dir = pathlib.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_file_bkp = "pipeline.config" if isinstance(config_path, str): # directly provide a config object. this usually used # when you want to train with several different parameters in # one script. config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) else: config = config_path proto_str = text_format.MessageToString(config, indent=2) with (model_dir / config_file_bkp).open("w") as f: f.write(proto_str) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config net = build_network(model_cfg).to(device) if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) target_assigner = net.target_assigner voxel_generator = net.voxel_generator class_names = target_assigner.classes # 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 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 center_limit_range = model_cfg.post_center_limit_range """ 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) dataloader = torch.utils.data.DataLoader( dataset, batch_size=input_cfg.batch_size, shuffle=True, num_workers=input_cfg.preprocess.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.preprocess.num_workers, pin_memory=False, collate_fn=merge_second_batch) data_iter = iter(dataloader) print(data_iter) ###################### # TRAINING ###################### model_logging = SimpleModelLog(model_dir) model_logging.open() model_logging.log_text(proto_str + "\n", 0, tag="config") total_step_elapsed = 0 remain_steps = train_cfg.steps - net.get_global_step() t = time.time() ckpt_start_time = t steps_per_eval = train_cfg.steps_per_eval total_loop = train_cfg.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.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] ret_dict = net(example_torch) # FCOS losses = ret_dict['total_loss'] loss_cls = ret_dict["loss_cls"] loss_reg = ret_dict["loss_reg"] cls_preds = ret_dict['cls_preds'] labels = ret_dict["labels"] cared = ret_dict["labels"] optimizer.zero_grad() losses.backward() #torch.nn.utils.clip_grad_norm_(net.parameters(), 1) # optimizer_step is for updating the parameter, so clip before update optimizer.step() net.update_global_step() #need to unpack the [0] for fpn net_metrics = net.update_metrics(loss_cls, loss_reg, cls_preds[0], labels, cared) step_time = (time.time() - t) t = time.time() metrics = {} global_step = net.get_global_step() #print log if global_step % display_step == 0: metrics["runtime"] = { "step": global_step, "steptime": step_time, } metrics.update(net_metrics) metrics["misc"] = { "num_vox": int(example_torch["voxels"].shape[0]), "lr": float(optimizer.lr), } model_logging.log_metrics(metrics, global_step) ckpt_elasped_time = time.time() - ckpt_start_time torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step()) 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) model_logging.log_text("#################################", global_step) model_logging.log_text("# EVAL", global_step) model_logging.log_text("#################################", global_step) model_logging.log_text("Generate output labels...", global_step) t = time.time() detections = [] prog_bar = ProgressBar() net.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) with torch.no_grad(): detections += net(example) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) model_logging.log_text( f'generate label finished({sec_per_ex:.2f}/s). start eval:', global_step) result_dict = eval_dataset.dataset.evaluation( detections, str(result_path_step)) for k, v in result_dict["results"].items(): model_logging.log_text("Evaluation {}".format(k), global_step) model_logging.log_text(v, global_step) model_logging.log_metrics(result_dict["detail"], global_step) with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(detections, f) net.train() ''' new version of evaluation while trainging # do the evaluation while traingingi if global_step % steps_per_eval == 0: 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) model_logging.log_text("#################################", global_step) model_logging.log_text("# EVAL", global_step) model_logging.log_text("#################################", global_step) model_logging.log_text("Generate output labels...", global_step) t = time.time() detections = [] prog_bar = ProgressBar() net.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) with torch.no_grad(): detections += net(example) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) model_logging.log_text( f'generate label finished({sec_per_ex:.2f}/s). start eval:', global_step) result_dict = eval_dataset.dataset.evaluation( detections, str(result_path_step)) for k, v in result_dict["results"].items(): model_logging.log_text("Evaluation {}".format(k), global_step) model_logging.log_text(v, global_step) model_logging.log_metrics(result_dict["detail"], global_step) with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(detections, f) net.train() ''' except Exception as e: print("trainging error") raise e finally: model_logging.close() # save model before exit torchplus.train.save_models(model_dir, [net, optimizer], net.get_global_step())
def train( config_path: Union[str, Path, pipeline.TrainEvalPipelineConfig], model_dir: Union[str, Path], data_root_path: Union[str, Path], result_path: Optional[Union[str, Path]] = None, display_step: int = 50, pretrained_path=None, pretrained_include=None, pretrained_exclude=None, freeze_include=None, freeze_exclude=None, measure_time: bool = False, resume: bool = False, ): """train a VoxelNet model specified by a config file. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_dir = real_path(model_dir, check_exists=False) 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) model_dir = Path(model_dir) if result_path is None: result_path = model_dir / "results" else: result_path = assert_real_path(result_path, mkdir=True) config_file_bkp = DEFAULT_CONFIG_FILE_NAME if isinstance(config_path, pipeline.TrainEvalPipelineConfig): # directly provide a config object. this usually used # when you want to train with several different parameters in # one script. config = config_path proto_str = text_format.MessageToString(config, use_short_repeated_primitives=True, indent=2) else: config_path = assert_real_path(config_path) data_root_path = assert_real_path(data_root_path) config = read_pipeline_config(config_path, data_root_path) # Copy the contents of config_path to config_file_bkp verbatim without passing it through the protobuf parser. with open(str(config_path), "r") as f: proto_str = f.read() with (model_dir / config_file_bkp).open("w") as f: f.write(proto_str) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config net = build_network(model_cfg, measure_time).to(device) if train_cfg.enable_mixed_precision: # net.half() net.metrics_to_float() net.convert_norm_to_float(net) target_assigner = net.target_assigner voxel_generator = net.voxel_generator # print("num parameters:", len(list(net.parameters()))) print("num parameters (million): ", count_parameters(net) * 1e-6) torchplus.train.try_restore_latest_checkpoints(model_dir, [net]) if pretrained_path is not None: model_dict = net.state_dict() pretrained_dict = torch.load(pretrained_path) pretrained_dict = filter_param_dict(pretrained_dict, pretrained_include, pretrained_exclude) new_pretrained_dict = {} for k, v in pretrained_dict.items(): if k in model_dict and v.shape == model_dict[k].shape: new_pretrained_dict[k] = v print("Load pretrained parameters:") for k, v in new_pretrained_dict.items(): print(k, v.shape) model_dict.update(new_pretrained_dict) net.load_state_dict(model_dict) freeze_params_v2(dict(net.named_parameters()), freeze_include, freeze_exclude) net.clear_global_step() net.clear_metrics() optimizer_cfg = train_cfg.optimizer loss_scale = train_cfg.loss_scale_factor fastai_optimizer = optimizer_builder.build( optimizer_cfg, net, mixed=False, loss_scale=loss_scale) if loss_scale < 0: loss_scale = "dynamic" amp_optimizer = fastai_optimizer torchplus.train.try_restore_latest_checkpoints(model_dir,[amp_optimizer]) float_dtype = torch.float32 collate_fn = merge_second_batch num_gpu = 1 ###################### # PREPARE INPUT ###################### def get_train_dataloader(input_cfg, model_cfg, voxel_generator, target_assigner, multi_gpu, num_gpu, collate_fn, _worker_init_fn): 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=True, collate_fn=collate_fn, worker_init_fn=_worker_init_fn, drop_last=not multi_gpu) return dataloader eval_dataset = input_reader_builder.build( eval_input_cfg, model_cfg, training=False, voxel_generator=voxel_generator, target_assigner=target_assigner) eval_dataloader = torch.utils.data.DataLoader( eval_dataset, batch_size=eval_input_cfg.batch_size, # only support multi-gpu train shuffle=False, num_workers=eval_input_cfg.preprocess.num_workers, pin_memory=False, collate_fn=merge_second_batch) ###################### # TRAINING ###################### model_logging = SimpleModelLog(model_dir) model_logging.open() model_logging.log_text(proto_str + "\n", 0, tag="config") epochs = train_cfg.steps epochs_per_eval = train_cfg.steps_per_eval clear_metrics_every_epoch = train_cfg.clear_metrics_every_epoch amp_optimizer.zero_grad() step_times = [] eval_times = [] t = time.time() reset_ds_epoch = False run_once = True if not (os.getenv("MLFLOW_EXPERIMENT_ID") or os.getenv("MLFLOW_EXPERIMENT_NAME")): mlflow.set_experiment("object_detection") try: while True: if run_once or reset_ds_epoch: dataloader = get_train_dataloader(input_cfg, model_cfg, voxel_generator, target_assigner, multi_gpu, num_gpu, collate_fn, _worker_init_fn) total_step = int(np.ceil((len(dataloader.dataset) / dataloader.batch_size) * epochs)) steps_per_eval = int(np.floor((len(dataloader.dataset) / dataloader.batch_size) * epochs_per_eval)) train_cfg.steps = int(total_step) train_cfg.steps_per_eval = int(steps_per_eval) lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, amp_optimizer, total_step) print(f"\nnumber of samples: {len(dataloader.dataset)}\ntotal_steps: {total_step}\nsteps_per_eval: {steps_per_eval}") run_once = False if clear_metrics_every_epoch: net.clear_metrics() for example in dataloader: lr_scheduler.step(net.get_global_step()) time_metrics = example["metrics"] example.pop("metrics") example_torch = example_convert_to_torch(example, float_dtype) batch_size = example["anchors"].shape[0] ret_dict = net(example_torch) 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(), 30.0) # torch.nn.utils.clip_grad_norm_(amp.master_params(amp_optimizer), 10.0) amp_optimizer.step() amp_optimizer.zero_grad() net.update_global_step() global_step = net.get_global_step() net_metrics = net.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()) if global_step % display_step == 0: if measure_time: for name, val in net.get_avg_time_dict().items(): print(f"avg {name} time = {val * 1000:.3f} ms") loc_loss_elem = [ float(loc_loss[:, :, i].sum().detach().cpu().numpy() / batch_size) for i in range(loc_loss.shape[-1]) ] total_seconds = ((total_step - global_step) * np.mean(step_times)) if len(eval_times) != 0: eval_seconds = ((epochs / epochs_per_eval) - len(eval_times)) * np.mean(eval_times) total_seconds += eval_seconds next_eval_seconds = (steps_per_eval - (global_step % steps_per_eval)) * np.mean(step_times) metrics["runtime"] = { "step": global_step, "steptime": np.mean(step_times), "ETA": seconds_to_eta(total_seconds), "eval_ETA": seconds_to_eta(next_eval_seconds), } 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) # if global_step % steps_per_eval != 0 and global_step % 1000 == 0: # torchplus.train.save_models(model_dir, [net, amp_optimizer], net.get_global_step()) if global_step % steps_per_eval == 0: torchplus.train.save_models(model_dir, [net, amp_optimizer], global_step) net.eval() result_path_step = result_path / f"step_{global_step}" result_path_step.mkdir(parents=True, exist_ok=True) model_logging.log_text("#################################", global_step) model_logging.log_text("# EVAL", global_step) model_logging.log_text("#################################", global_step) model_logging.log_text("Generate output labels...", global_step) t = time.time() detections = [] prog_bar = ProgressBar() net.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) detections += net(example) prog_bar.print_bar() sec_per_ex = len(eval_dataset) / (time.time() - t) eval_times.append((time.time() - t)) model_logging.log_text(f'generate label finished({sec_per_ex:.2f}/s). start eval:', global_step) result_dict = eval_dataset.dataset.evaluation(detections, result_path_step) if result_dict is None: raise RuntimeError("eval_dataset.dataset.evaluation() returned None") for k, v in result_dict["results"].items(): model_logging.log_text("Evaluation {}".format(k), global_step) model_logging.log_text(v, global_step) model_logging.log_metrics(result_dict["detail"], global_step) with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(detections, f) net.train() if global_step >= total_step: break if net.get_global_step() >= total_step: break except Exception as e: if 'example' in locals(): print(json.dumps(example["metadata"], indent=2)) global_step = net.get_global_step() model_logging.log_text(str(e), global_step) if 'example' in locals(): model_logging.log_text(json.dumps(example["metadata"], indent=2), global_step) torchplus.train.save_models(model_dir, [net, amp_optimizer], global_step) raise e finally: model_logging.close() torchplus.train.save_models(model_dir, [net, amp_optimizer], net.get_global_step()) def _save_checkpoint_info(file_path, config_filename, checkpoint_filename): from yaml import dump with open(file_path, "w") as config_info_file: checkpoint_info = { "config": config_filename, "checkpoint": checkpoint_filename } dump(checkpoint_info, config_info_file, default_flow_style=False) ckpt_info_path = str(model_dir / "checkpoint_info.yaml") latest_ckpt_filename = "voxelnet-{}.tckpt".format(net.get_global_step()) _save_checkpoint_info(ckpt_info_path, config_file_bkp, latest_ckpt_filename) mlflow.log_artifact(ckpt_info_path, "model") mlflow.log_artifact(str(model_dir / config_file_bkp), "model") mlflow.log_artifact(str(model_dir / latest_ckpt_filename), "model")
else: net_parallel = net optimizer_cfg = train_cfg.optimizer loss_scale = train_cfg.loss_scale_factor fastai_optimizer = optimizer_builder.build(optimizer_cfg, net, mixed=False, loss_scale=loss_scale) if loss_scale < 0: loss_scale = "dynamic" amp_optimizer = fastai_optimizer lr_scheduler = lr_scheduler_builder.build(optimizer_cfg, amp_optimizer, train_cfg.steps) float_dtype = torch.float32 if cfg.multi_gpu: num_gpu = torch.cuda.device_count() print(f"MULTI-GPU: use {num_gpu} gpu") collate_fn = merge_second_batch_multigpu else: collate_fn = merge_second_batch num_gpu = 1 ###################### # PREPARE INPUT ###################### dataset = input_reader_builder.build(input_cfg,
def train(config_path, model_dir, result_path=None, create_folder=False, display_step=50, summary_step=5, pickle_result=True, resume=False): """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) 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_file_bkp = "pipeline.config" if isinstance(config_path, str): # directly provide a config object. this usually used # when you want to train with several different parameters in # one script. config = pipeline_pb2.TrainEvalPipelineConfig() with open(config_path, "r") as f: proto_str = f.read() text_format.Merge(proto_str, config) else: config = config_path proto_str = text_format.MessageToString(config, indent=2) with (model_dir / config_file_bkp).open("w") as f: f.write(proto_str) input_cfg = config.train_input_reader eval_input_cfg = config.eval_input_reader model_cfg = config.model.second train_cfg = config.train_config net = build_network(model_cfg).cuda() if train_cfg.enable_mixed_precision: net.half() net.metrics_to_float() net.convert_norm_to_float(net) target_assigner = net.target_assigner voxel_generator = net.voxel_generator class_names = target_assigner.classes # 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 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 center_limit_range = model_cfg.post_center_limit_range """ 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) 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 ###################### 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 # 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.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] 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["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) if "image_info" in example['metadata'][0]: metrics["image_idx"] = example['metadata'][0][ "image_info"]['image_idx'] training_detail.append(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) """ log_str = metric_to_str(metrics) 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()) 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) for example in iter(eval_dataloader): example = example_convert_to_torch(example, float_dtype) dt_annos += predict_to_kitti_label(net, example, 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) result_official, result_coco = eval_dataset.dataset.evaluation( dt_annos) print(result_official) print(result_official, file=logf) print(result_coco) print(result_coco, file=logf) if pickle_result: with open(result_path_step / "result.pkl", 'wb') as f: pickle.dump(dt_annos, f) else: kitti_anno_to_label_file(dt_annos, result_path_step) writer.add_text('eval_result', result_official, global_step) writer.add_text('eval_result coco', result_coco, 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(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()