if __name__ == '__main__': configs = parse_eval_configs() configs.distributed = False # For evaluation class_names = load_classes(configs.classnames_infor_path) model = create_model(configs) model.print_network() print('\n\n' + '-*=' * 30 + '\n\n') assert os.path.isfile(configs.pretrained_path), "No file at {}".format( configs.pretrained_path) # model.load_weights(configs.pretrained_path) model.load_state_dict(torch.load(configs.pretrained_path)) configs.device = torch.device( 'cpu' if configs.no_cuda else 'cuda:{}'.format(configs.gpu_idx)) model = model.to(device=configs.device) model.eval() print('Create the validation dataloader') val_dataloader = create_val_dataloader(configs) print("\nStart computing mAP...\n") precision, recall, AP, f1, ap_class = evaluate_mAP(val_dataloader, model, configs, None) print("\nDone computing mAP...\n") for idx, cls in enumerate(ap_class): print("\t>>>\t Class {} ({}): precision = {:.4f}, recall = {:.4f}, AP = {:.4f}, f1: {:.4f}".format(cls, \ class_names[cls][:3], precision[idx], recall[idx], AP[idx], f1[idx])) print("\nmAP: {}\n".format(AP.mean()))
max_voxels=20000) model = create_model(configs, voxel_generator) print('\n\n' + '-*=' * 30 + '\n\n') assert os.path.isfile(configs.pretrained_path), "No file at {}".format( configs.pretrained_path) model.load_state_dict( torch.load(configs.pretrained_path, map_location='cpu')) print('Loaded weights from {}\n'.format(configs.pretrained_path)) configs.device = torch.device( 'cpu' if configs.no_cuda else 'cuda:{}'.format(configs.gpu_idx)) model = model.to(device=configs.device) out_cap = None model.eval() print('Create the validation dataloader') val_dataloader = create_val_dataloader(configs, voxel_generator) print("\nStart computing mAP...\n") precision, recall, AP, f1, ap_class = evaluate_mAP(val_dataloader, model, configs, None) print("\nDone computing mAP...\n") for idx, cls in enumerate(ap_class): print("\t>>>\t Class {} ({}): precision = {:.4f}, recall = {:.4f}, AP = {:.4f}, f1: {:.4f}".format(cls, \ class_names[cls][:3], precision[idx], recall[idx], AP[idx], f1[idx])) print("\nmAP: {}\n".format(AP.mean()))
def main_worker(gpu_idx, configs): configs.gpu_idx = gpu_idx if configs.gpu_idx is not None: print("Use GPU: {} for training".format(configs.gpu_idx)) configs.device = torch.device('cuda:{}'.format(configs.gpu_idx)) if configs.distributed: if configs.dist_url == "env://" and configs.rank == -1: configs.rank = int(os.environ["RANK"]) if configs.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx dist.init_process_group(backend=configs.dist_backend, init_method=configs.dist_url, world_size=configs.world_size, rank=configs.rank) configs.is_master_node = (not configs.distributed) or ( configs.distributed and (configs.rank % configs.ngpus_per_node == 0)) if configs.is_master_node: logger = Logger(configs.logs_dir, configs.saved_fn) logger.info('>>> Created a new logger') logger.info('>>> configs: {}'.format(configs)) tb_writer = SummaryWriter(log_dir=os.path.join(configs.logs_dir, 'tensorboard')) else: logger = None tb_writer = None # model model = create_model(configs) # load weight from a checkpoint if configs.pretrained_path is not None: assert os.path.isfile(configs.pretrained_path), "=> no checkpoint found at '{}'".format(configs.pretrained_path) model.load_state_dict(torch.load(configs.pretrained_path)) if logger is not None: logger.info('loaded pretrained model at {}'.format(configs.pretrained_path)) # resume weights of model from a checkpoint if configs.resume_path is not None: assert os.path.isfile(configs.resume_path), "=> no checkpoint found at '{}'".format(configs.resume_path) model.load_state_dict(torch.load(configs.pretrained_path)) if logger is not None: logger.info('resume training model from checkpoint {}'.format(configs.pretrained_path)) # Data Parallel model = make_data_parallel(model, configs) # Make sure to create optimizer after moving the model to cuda optimizer = create_optimizer(configs, model) lr_scheduler = create_lr_scheduler(optimizer, configs) # resume optimizer, lr_scheduler from a checkpoint if configs.resume_path is not None: utils_path = configs.resume_path.replace('Model_', 'Utils_') assert os.path.isfile(utils_path), "=> no checkpoint found at '{}'".format(utils_path) utils_state_dict = torch.load(utils_path, map_location='cuda:{}'.format(configs.gpu_idx)) optimizer.load_state_dict(utils_state_dict['optimizer']) lr_scheduler.load_state_dict(utils_state_dict['lr_scheduler']) configs.start_epoch = utils_state_dict['epoch'] + 1 if configs.is_master_node: num_parameters = get_num_parameters(model) logger.info('number of trained parameters of the model: {}'.format(num_parameters)) if logger is not None: logger.info(">>> Loading dataset & getting dataloader...") # Create dataloader train_dataloader, train_sampler = create_train_dataloader(configs) if logger is not None: logger.info('number of batches in training set: {}'.format(len(train_dataloader))) if configs.evaluate: val_dataloader = create_val_dataloader(configs) precision, recall, AP, f1, ap_class = evaluate_mAP(val_dataloader, model, configs, None) print('Evaluate - precision: {}, recall: {}, AP: {}, f1: {}, ap_class: {}'.format(precision, recall, AP, f1, ap_class)) return for epoch in range(configs.start_epoch, configs.num_epochs + 1): if logger is not None: logger.info('{}'.format('*-' * 40)) logger.info('{} {}/{} {}'.format('=' * 35, epoch, configs.num_epochs, '=' * 35)) logger.info('{}'.format('*-' * 40)) logger.info('>>> Epoch: [{}/{}]'.format(epoch, configs.num_epochs)) if configs.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_one_epoch(train_dataloader, model, optimizer, lr_scheduler, epoch, configs, logger, tb_writer) if not configs.no_val: val_dataloader = create_val_dataloader(configs) precision, recall, AP, f1, ap_class = evaluate_mAP(val_dataloader, model, configs, logger) val_metrics_dict = {'precision': precision, 'recall': recall, 'AP': AP, 'f1': f1, 'ap_class': ap_class} if tb_writer is not None: tb_writer.add_scalars('Validation', val_metrics_dict, epoch) # Save checkpoint if configs.is_master_node and ((epoch % configs.checkpoint_freq) == 0): model_state_dict, utils_state_dict = get_saved_state(model, optimizer, lr_scheduler, epoch, configs) save_checkpoint(configs.checkpoints_dir, configs.saved_fn, model_state_dict, utils_state_dict, epoch) if tb_writer is not None: tb_writer.close() if configs.distributed: cleanup()
def main_worker(gpu_idx, configs): configs.gpu_idx = gpu_idx configs.device = torch.device('cpu' if configs.gpu_idx is None else 'cuda:{}'.format(configs.gpu_idx)) voxel_generator = VoxelGeneratorV2( voxel_size=list(configs.voxel_size), point_cloud_range=list(configs.point_cloud_range), max_num_points=configs.max_number_of_points_per_voxel, max_voxels=20000) if configs.distributed: if configs.dist_url == "env://" and configs.rank == -1: configs.rank = int(os.environ["RANK"]) if configs.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx dist.init_process_group(backend=configs.dist_backend, init_method=configs.dist_url, world_size=configs.world_size, rank=configs.rank) configs.subdivisions = int(64 / configs.batch_size / configs.ngpus_per_node) else: configs.subdivisions = int(64 / configs.batch_size) configs.is_master_node = (not configs.distributed) or ( configs.distributed and (configs.rank % configs.ngpus_per_node == 0)) if configs.is_master_node: logger = Logger(configs.logs_dir, configs.saved_fn) logger.info('>>> Created a new logger') logger.info('>>> configs: {}'.format(configs)) tb_writer = SummaryWriter( log_dir=os.path.join(configs.logs_dir, 'tensorboard')) else: logger = None tb_writer = None # model model = create_model(configs, voxel_generator) input = torch.rand(( 1999, 100, 4, ), device='cpu') input1 = torch.rand((1999, 4), device='cpu') input2 = torch.rand(1999, device='cpu') #if tb_writer is not None: # tb_writer.add_graph(model, (input, input1, input2)) # load weight from a checkpoint if configs.pretrained_path is not None: assert os.path.isfile( configs.pretrained_path), "=> no checkpoint found at '{}'".format( configs.pretrained_path) model.load_state_dict( torch.load(configs.pretrained_path, map_location='cpu')) if logger is not None: logger.info('loaded pretrained model at {}'.format( configs.pretrained_path)) # resume weights of model from a checkpoint if configs.resume_path is not None: assert os.path.isfile( configs.resume_path), "=> no checkpoint found at '{}'".format( configs.resume_path) model.load_state_dict( torch.load(configs.resume_path, map_location='cpu')) if logger is not None: logger.info('resume training model from checkpoint {}'.format( configs.resume_path)) # Data Parallel model = make_data_parallel(model, configs) # Make sure to create optimizer after moving the model to cuda optimizer = create_optimizer(configs, model) lr_scheduler = create_lr_scheduler(optimizer, configs) configs.step_lr_in_epoch = False if configs.lr_type in [ 'multi_step', 'cosin', 'one_cycle' ] else True # resume optimizer, lr_scheduler from a checkpoint if configs.resume_path is not None: utils_path = configs.resume_path.replace('Model_', 'Utils_') assert os.path.isfile( utils_path), "=> no checkpoint found at '{}'".format(utils_path) utils_state_dict = torch.load(utils_path, map_location='cuda:{}'.format( configs.gpu_idx)) optimizer.load_state_dict(utils_state_dict['optimizer']) lr_scheduler.load_state_dict(utils_state_dict['lr_scheduler']) configs.start_epoch = utils_state_dict['epoch'] + 1 if configs.is_master_node: num_parameters = get_num_parameters(model) logger.info('number of trained parameters of the model: {}'.format( num_parameters)) if logger is not None: logger.info(">>> Loading dataset & getting dataloader...") # Create dataloader train_dataloader, train_sampler = create_train_dataloader( configs, voxel_generator) if logger is not None: logger.info('number of batches in training set: {}'.format( len(train_dataloader))) if configs.evaluate: val_dataloader = create_val_dataloader(configs, voxel_generator) val_loss = validate(val_dataloader, model, configs) print('val_loss: {:.4e}'.format(val_loss)) return for epoch in range(configs.start_epoch, configs.num_epochs + 1): if logger is not None: logger.info('{}'.format('*-' * 40)) logger.info('{} {}/{} {}'.format('=' * 35, epoch, configs.num_epochs, '=' * 35)) logger.info('{}'.format('*-' * 40)) logger.info('>>> Epoch: [{}/{}]'.format(epoch, configs.num_epochs)) if configs.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_one_epoch(train_dataloader, model, optimizer, lr_scheduler, epoch, configs, logger, tb_writer) if (not configs.no_val) and (epoch % configs.checkpoint_freq == 0): val_dataloader = create_val_dataloader(configs, voxel_generator) print('number of batches in val_dataloader: {}'.format( len(val_dataloader))) val_loss = validate(val_dataloader, model, configs) print('val_loss: {:.4e}'.format(val_loss)) if tb_writer is not None: tb_writer.add_scalar('Val_loss', val_loss, epoch) # Save checkpoint if configs.is_master_node and ((epoch % configs.checkpoint_freq) == 0): model_state_dict, utils_state_dict = get_saved_state( model, optimizer, lr_scheduler, epoch, configs) save_checkpoint(configs.checkpoints_dir, configs.saved_fn, model_state_dict, utils_state_dict, epoch) if not configs.step_lr_in_epoch: lr_scheduler.step() if tb_writer is not None: tb_writer.add_scalar('LR', lr_scheduler.get_lr()[0], epoch) if tb_writer is not None: tb_writer.close() if configs.distributed: cleanup()