def load_pretrained_model(self): if type(self.initialize_from) != list: self.initialize_from = [self.initialize_from] self.modules_to_initialize = [self.modules_to_initialize] for init_from, modules_to_init in zip(self.initialize_from, self.modules_to_initialize): print(init_from) checkpoint = get_last_checkpoint(init_from) checkpoint = torch.load(checkpoint) new_state_dict = self.state_dict() for key in checkpoint['state_dict'].keys(): for k in key.split('.'): if k in modules_to_init: new_state_dict[key] = checkpoint['state_dict'][key] print('pretrain parameters: %s' % k) self.load_state_dict(new_state_dict)
def test(cfg): """ Test a model """ logging.setup_logging(logger, cfg) logger.info("Test with config") logger.info(pprint.pformat(cfg)) model = model_builder.build_model(cfg) if du.is_master_proc(): misc.log_model_info(model) if cfg.TEST.CHECKPOINT_FILE_PATH != "": logger.info("Load from given checkpoint file.") gs, checkpoint_epoch = cu.load_checkpoint( cfg.TEST.CHECKPOINT_FILE_PATH, model, cfg.NUM_GPUS > 1, optimizer=None, inflation=False, convert_from_caffe2=False) start_epoch = checkpoint_epoch + 1 elif cfg.TRAIN.AUTO_RESUME and cu.has_checkpoint(cfg.OUTPUT_DIR): logger.info("Load from last checkpoint.") last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR) gs, checkpoint_epoch = cu.load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1, None) start_epoch = checkpoint_epoch + 1 # Create the video train and val loaders. test_loader = loader.construct_loader(cfg, "test") test_meter = TestMeter(cfg) if cfg.TEST.AUGMENT_TEST: evaluate_with_augmentation(test_loader, model, test_meter, cfg) else: evaluate(test_loader, model, test_meter, cfg)
def csgd_train_main(local_rank, cfg: BaseConfigByEpoch, target_deps, succeeding_strategy, pacesetter_dict, centri_strength, pruned_weights, net=None, train_dataloader=None, val_dataloader=None, show_variables=False, convbuilder=None, init_hdf5=None, no_l2_keywords='depth', use_nesterov=False, load_weights_keyword=None, keyword_to_lr_mult=None, auto_continue=False, save_hdf5_epochs=10000): ensure_dir(cfg.output_dir) ensure_dir(cfg.tb_dir) clusters_save_path = os.path.join(cfg.output_dir, 'clusters.npy') with Engine(local_rank=local_rank) as engine: engine.setup_log(name='train', log_dir=cfg.output_dir, file_name='log.txt') # ----------------------------- build model ------------------------------ if convbuilder is None: convbuilder = ConvBuilder(base_config=cfg) if net is None: net_fn = get_model_fn(cfg.dataset_name, cfg.network_type) model = net_fn(cfg, convbuilder) else: model = net model = model.cuda() # ----------------------------- model done ------------------------------ # ---------------------------- prepare data ------------------------- if train_dataloader is None: train_data = create_dataset(cfg.dataset_name, cfg.dataset_subset, cfg.global_batch_size, distributed=engine.distributed) if cfg.val_epoch_period > 0 and val_dataloader is None: val_data = create_dataset(cfg.dataset_name, 'val', global_batch_size=100, distributed=False) engine.echo('NOTE: Data prepared') engine.echo( 'NOTE: We have global_batch_size={} on {} GPUs, the allocated GPU memory is {}' .format(cfg.global_batch_size, torch.cuda.device_count(), torch.cuda.memory_allocated())) # ----------------------------- data done -------------------------------- # ------------------------ parepare optimizer, scheduler, criterion ------- if no_l2_keywords is None: no_l2_keywords = [] if type(no_l2_keywords) is not list: no_l2_keywords = [no_l2_keywords] # For a target parameter, cancel its weight decay in optimizer, because the weight decay will be later encoded in the decay mat conv_idx = 0 for k, v in model.named_parameters(): if v.dim() != 4: continue print('prune {} from {} to {}'.format(conv_idx, target_deps[conv_idx], cfg.deps[conv_idx])) if target_deps[conv_idx] < cfg.deps[conv_idx]: no_l2_keywords.append(k.replace(KERNEL_KEYWORD, 'conv')) no_l2_keywords.append(k.replace(KERNEL_KEYWORD, 'bn')) conv_idx += 1 print('no l2: ', no_l2_keywords) optimizer = get_optimizer(engine, cfg, model, no_l2_keywords=no_l2_keywords, use_nesterov=use_nesterov, keyword_to_lr_mult=keyword_to_lr_mult) scheduler = get_lr_scheduler(cfg, optimizer) criterion = get_criterion(cfg).cuda() # --------------------------------- done ------------------------------- engine.register_state(scheduler=scheduler, model=model, optimizer=optimizer) if engine.distributed: torch.cuda.set_device(local_rank) engine.echo('Distributed training, device {}'.format(local_rank)) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], broadcast_buffers=False, ) else: assert torch.cuda.device_count() == 1 engine.echo('Single GPU training') if cfg.init_weights: engine.load_checkpoint(cfg.init_weights) if init_hdf5: engine.load_hdf5(init_hdf5, load_weights_keyword=load_weights_keyword) if auto_continue: assert cfg.init_weights is None engine.load_checkpoint(get_last_checkpoint(cfg.output_dir)) if show_variables: engine.show_variables() # ===================================== prepare the clusters and matrices for C-SGD ========== kernel_namedvalue_list = engine.get_all_conv_kernel_namedvalue_as_list( ) if os.path.exists(clusters_save_path): layer_idx_to_clusters = np.load(clusters_save_path, allow_pickle=True).item() else: if local_rank == 0: layer_idx_to_clusters = get_layer_idx_to_clusters( kernel_namedvalue_list=kernel_namedvalue_list, target_deps=target_deps, pacesetter_dict=pacesetter_dict) if pacesetter_dict is not None: for follower_idx, pacesetter_idx in pacesetter_dict.items( ): if pacesetter_idx in layer_idx_to_clusters: layer_idx_to_clusters[ follower_idx] = layer_idx_to_clusters[ pacesetter_idx] np.save(clusters_save_path, layer_idx_to_clusters) else: while not os.path.exists(clusters_save_path): time.sleep(10) print('sleep, waiting for process 0 to calculate clusters') layer_idx_to_clusters = np.load(clusters_save_path, allow_pickle=True).item() param_name_to_merge_matrix = generate_merge_matrix_for_kernel( deps=cfg.deps, layer_idx_to_clusters=layer_idx_to_clusters, kernel_namedvalue_list=kernel_namedvalue_list) add_vecs_to_merge_mat_dicts(param_name_to_merge_matrix) param_name_to_decay_matrix = generate_decay_matrix_for_kernel_and_vecs( deps=cfg.deps, layer_idx_to_clusters=layer_idx_to_clusters, kernel_namedvalue_list=kernel_namedvalue_list, weight_decay=cfg.weight_decay, weight_decay_bias=cfg.weight_decay_bias, centri_strength=centri_strength) print(param_name_to_decay_matrix.keys()) print(param_name_to_merge_matrix.keys()) conv_idx = 0 param_to_clusters = {} for k, v in model.named_parameters(): if v.dim() != 4: continue if conv_idx in layer_idx_to_clusters: for clsts in layer_idx_to_clusters[conv_idx]: if len(clsts) > 1: param_to_clusters[v] = layer_idx_to_clusters[conv_idx] break conv_idx += 1 # ============================================================================================ # ------------ do training ---------------------------- # engine.log("\n\nStart training with pytorch version {}".format( torch.__version__)) iteration = engine.state.iteration iters_per_epoch = num_iters_per_epoch(cfg) max_iters = iters_per_epoch * cfg.max_epochs tb_writer = SummaryWriter(cfg.tb_dir) tb_tags = ['Top1-Acc', 'Top5-Acc', 'Loss'] model.train() done_epochs = iteration // iters_per_epoch last_epoch_done_iters = iteration % iters_per_epoch if done_epochs == 0 and last_epoch_done_iters == 0: engine.save_hdf5(os.path.join(cfg.output_dir, 'init.hdf5')) recorded_train_time = 0 recorded_train_examples = 0 collected_train_loss_sum = 0 collected_train_loss_count = 0 for epoch in range(done_epochs, cfg.max_epochs): if engine.distributed and hasattr(train_data, 'train_sampler'): train_data.train_sampler.set_epoch(epoch) if epoch == done_epochs: pbar = tqdm(range(iters_per_epoch - last_epoch_done_iters)) else: pbar = tqdm(range(iters_per_epoch)) if epoch == 0 and local_rank == 0: val_during_train(epoch=epoch, iteration=iteration, tb_tags=tb_tags, engine=engine, model=model, val_data=val_data, criterion=criterion, descrip_str='Init', dataset_name=cfg.dataset_name, test_batch_size=TEST_BATCH_SIZE, tb_writer=tb_writer) top1 = AvgMeter() top5 = AvgMeter() losses = AvgMeter() discrip_str = 'Epoch-{}/{}'.format(epoch, cfg.max_epochs) pbar.set_description('Train' + discrip_str) for _ in pbar: start_time = time.time() data, label = load_cuda_data(train_data, dataset_name=cfg.dataset_name) # load_cuda_data(train_dataloader, cfg.dataset_name) data_time = time.time() - start_time train_net_time_start = time.time() acc, acc5, loss = train_one_step( model, data, label, optimizer, criterion, param_name_to_merge_matrix=param_name_to_merge_matrix, param_name_to_decay_matrix=param_name_to_decay_matrix) train_net_time_end = time.time() if iteration > TRAIN_SPEED_START * max_iters and iteration < TRAIN_SPEED_END * max_iters: recorded_train_examples += cfg.global_batch_size recorded_train_time += train_net_time_end - train_net_time_start scheduler.step() for module in model.modules(): if hasattr(module, 'set_cur_iter'): module.set_cur_iter(iteration) if iteration % cfg.tb_iter_period == 0 and engine.world_rank == 0: for tag, value in zip( tb_tags, [acc.item(), acc5.item(), loss.item()]): tb_writer.add_scalars(tag, {'Train': value}, iteration) deviation_sum = 0 for param, clusters in param_to_clusters.items(): pvalue = param.detach().cpu().numpy() for cl in clusters: if len(cl) == 1: continue selected = pvalue[cl, :, :, :] mean_kernel = np.mean(selected, axis=0, keepdims=True) diff = selected - mean_kernel deviation_sum += np.sum(diff**2) tb_writer.add_scalars('deviation_sum', {'Train': deviation_sum}, iteration) top1.update(acc.item()) top5.update(acc5.item()) losses.update(loss.item()) if epoch >= cfg.max_epochs - COLLECT_TRAIN_LOSS_EPOCHS: collected_train_loss_sum += loss.item() collected_train_loss_count += 1 pbar_dic = OrderedDict() pbar_dic['data-time'] = '{:.2f}'.format(data_time) pbar_dic['cur_iter'] = iteration pbar_dic['lr'] = scheduler.get_lr()[0] pbar_dic['top1'] = '{:.5f}'.format(top1.mean) pbar_dic['top5'] = '{:.5f}'.format(top5.mean) pbar_dic['loss'] = '{:.5f}'.format(losses.mean) pbar.set_postfix(pbar_dic) iteration += 1 if iteration >= max_iters or iteration % cfg.ckpt_iter_period == 0: engine.update_iteration(iteration) if (not engine.distributed) or (engine.distributed and engine.world_rank == 0): engine.save_and_link_checkpoint(cfg.output_dir) if iteration >= max_iters: break # do something after an epoch? engine.update_iteration(iteration) engine.save_latest_ckpt(cfg.output_dir) if (epoch + 1) % save_hdf5_epochs == 0: engine.save_hdf5( os.path.join(cfg.output_dir, 'epoch-{}.hdf5'.format(epoch))) if local_rank == 0 and \ cfg.val_epoch_period > 0 and (epoch >= cfg.max_epochs - 10 or epoch % cfg.val_epoch_period == 0): val_during_train(epoch=epoch, iteration=iteration, tb_tags=tb_tags, engine=engine, model=model, val_data=val_data, criterion=criterion, descrip_str=discrip_str, dataset_name=cfg.dataset_name, test_batch_size=TEST_BATCH_SIZE, tb_writer=tb_writer) if iteration >= max_iters: break # do something after the training if recorded_train_time > 0: exp_per_sec = recorded_train_examples / recorded_train_time else: exp_per_sec = 0 engine.log( 'TRAIN speed: from {} to {} iterations, batch_size={}, examples={}, total_net_time={:.4f}, examples/sec={}' .format(int(TRAIN_SPEED_START * max_iters), int(TRAIN_SPEED_END * max_iters), cfg.global_batch_size, recorded_train_examples, recorded_train_time, exp_per_sec)) if cfg.save_weights: engine.save_checkpoint(cfg.save_weights) print('NOTE: training finished, saved to {}'.format( cfg.save_weights)) engine.save_hdf5(os.path.join(cfg.output_dir, 'finish.hdf5')) if collected_train_loss_count > 0: engine.log( 'TRAIN LOSS collected over last {} epochs: {:.6f}'.format( COLLECT_TRAIN_LOSS_EPOCHS, collected_train_loss_sum / collected_train_loss_count)) if local_rank == 0: csgd_prune_and_save(engine=engine, layer_idx_to_clusters=layer_idx_to_clusters, save_file=pruned_weights, succeeding_strategy=succeeding_strategy, new_deps=target_deps)
def test(cfg): """ Perform multi-view testing on the pretrained video model. Args: cfg (CfgNode): configs. Details can be found in slowfast/config/defaults.py """ # Set up environment. du.init_distributed_training(cfg) # Set random seed from configs. np.random.seed(cfg.RNG_SEED) torch.manual_seed(cfg.RNG_SEED) # Setup logging format. logging.setup_logging(cfg.OUTPUT_DIR) # Print config. logger.info("Test with config:") logger.info(cfg) # Build the video model and print model statistics. model = build_model(cfg) if du.is_master_proc() and cfg.LOG_MODEL_INFO: misc.log_model_info(model, cfg, is_train=False) # Load a checkpoint to test if applicable. if cfg.TEST.CHECKPOINT_FILE_PATH != "": cu.load_checkpoint( cfg.TEST.CHECKPOINT_FILE_PATH, model, cfg.NUM_GPUS > 1, None, inflation=False, convert_from_caffe2=cfg.TEST.CHECKPOINT_TYPE == "caffe2", ) elif cu.has_checkpoint(cfg.OUTPUT_DIR): last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR) cu.load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1) elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "": # If no checkpoint found in TEST.CHECKPOINT_FILE_PATH or in the current # checkpoint folder, try to load checkpint from # TRAIN.CHECKPOINT_FILE_PATH and test it. cu.load_checkpoint( cfg.TRAIN.CHECKPOINT_FILE_PATH, model, cfg.NUM_GPUS > 1, None, inflation=False, convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2", ) else: # raise NotImplementedError("Unknown way to load checkpoint.") logger.info("Testing with random initialization. Only for debugging.") # Create video testing loaders. test_loader = loader.construct_loader(cfg, "test") logger.info("Testing model for {} iterations".format(len(test_loader))) if cfg.DETECTION.ENABLE: assert cfg.NUM_GPUS == cfg.TEST.BATCH_SIZE test_meter = AVAMeter(len(test_loader), cfg, mode="test") else: assert ( len(test_loader.dataset) % (cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS) == 0) # Create meters for multi-view testing. test_meter = TestMeter( len(test_loader.dataset) // (cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS), cfg.TEST.NUM_ENSEMBLE_VIEWS * cfg.TEST.NUM_SPATIAL_CROPS, cfg.MODEL.NUM_CLASSES, len(test_loader), cfg.DATA.MULTI_LABEL, cfg.DATA.ENSEMBLE_METHOD, ) # Set up writer for logging to Tensorboard format. if cfg.TENSORBOARD.ENABLE and du.is_master_proc( cfg.NUM_GPUS * cfg.NUM_SHARDS): writer = tb.TensorboardWriter(cfg) else: writer = None # # Perform multi-view test on the entire dataset. perform_test(test_loader, model, test_meter, cfg, writer) if writer is not None: writer.close()
def train(cfg): """ Train a video model for many epochs on train set and evaluate it on val set. Args: cfg (CfgNode): configs. Details can be found in slowfast/config/defaults.py """ # Setup logging format. logging.setup_logging(logger, cfg) # Print config. logger.info("Train with config:") logger.info(pprint.pformat(cfg)) # Build the video model and print model statistics. model = model_builder.build_model(cfg) if du.is_master_proc(): misc.log_model_info(model) # Construct the optimizer. optimizer = optim.construct_optimizer(model, cfg) # Record global step gs = 0 # Load a checkpoint to resume training if applicable. if cfg.TRAIN.AUTO_RESUME and cu.has_checkpoint(cfg.OUTPUT_DIR): logger.info("Load from last checkpoint.") last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR) gs, checkpoint_epoch = cu.load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1, optimizer) start_epoch = checkpoint_epoch + 1 elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "": logger.info("Load from given checkpoint file.") if cfg.TRAIN.LOAD_PART_OF_CHECKPOINT: gs, checkpoint_epoch = cu.load_part_of_checkpoint( cfg.TRAIN.CHECKPOINT_FILE_PATH, model, cfg.NUM_GPUS > 1, optimizer=None) else: gs, checkpoint_epoch = cu.load_checkpoint( cfg.TRAIN.CHECKPOINT_FILE_PATH, model, cfg.NUM_GPUS > 1, optimizer=None, inflation=False, convert_from_caffe2=False) start_epoch = checkpoint_epoch + 1 else: gs = 0 start_epoch = 0 # Create the video train and val loaders. train_loader = loader.construct_loader(cfg, "train") val_loader = loader.construct_loader(cfg, "val") # Create meters. train_meter = TrainMeter(len(train_loader), cfg) val_meter = ValMeter(cfg) # Perform the training loop. logger.info("Start epoch: {} gs {}".format(start_epoch + 1, gs + 1)) for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH): # Shuffle the dataset. loader.shuffle_dataset(train_loader, cur_epoch) # Evaluate the model on validation set. if misc.is_eval_epoch(cfg, cur_epoch): if cfg.TRAIN.USE_CENTER_VALIDATION: validation_epoch_center(val_loader, model, val_meter, cur_epoch, cfg) else: validation_epoch(val_loader, model, val_meter, cur_epoch, cfg) # Train for one epoch. gs = train_epoch(train_loader, model, optimizer, train_meter, cur_epoch, gs, cfg) # Compute precise BN stats. # if cfg.BN.USE_PRECISE_STATS and len(get_bn_modules(model)) > 0: # calculate_and_update_precise_bn( # train_loader, model, cfg.BN.NUM_BATCHES_PRECISE # ) # Save a checkpoint. if cu.is_checkpoint_epoch(cur_epoch, cfg.TRAIN.CHECKPOINT_PERIOD): cu.save_checkpoint(cfg.OUTPUT_DIR, model, optimizer, cur_epoch, gs, cfg)
def train_main(local_rank, cfg: BaseConfigByEpoch, net=None, train_dataloader=None, val_dataloader=None, show_variables=False, convbuilder=None, init_hdf5=None, no_l2_keywords='depth', gradient_mask=None, use_nesterov=False, tensorflow_style_init=False, load_weights_keyword=None, keyword_to_lr_mult=None, auto_continue=False, lasso_keyword_to_strength=None, save_hdf5_epochs=10000): if no_l2_keywords is None: no_l2_keywords = [] if type(no_l2_keywords) is not list: no_l2_keywords = [no_l2_keywords] ensure_dir(cfg.output_dir) ensure_dir(cfg.tb_dir) with Engine(local_rank=local_rank) as engine: engine.setup_log(name='train', log_dir=cfg.output_dir, file_name='log.txt') # ----------------------------- build model ------------------------------ if convbuilder is None: convbuilder = ConvBuilder(base_config=cfg) if net is None: net_fn = get_model_fn(cfg.dataset_name, cfg.network_type) model = net_fn(cfg, convbuilder) else: model = net model = model.cuda() # ----------------------------- model done ------------------------------ # ---------------------------- prepare data ------------------------- if train_dataloader is None: train_data = create_dataset(cfg.dataset_name, cfg.dataset_subset, cfg.global_batch_size, distributed=engine.distributed) if cfg.val_epoch_period > 0 and val_dataloader is None: val_data = create_dataset(cfg.dataset_name, 'val', global_batch_size=100, distributed=False) engine.echo('NOTE: Data prepared') engine.echo( 'NOTE: We have global_batch_size={} on {} GPUs, the allocated GPU memory is {}' .format(cfg.global_batch_size, torch.cuda.device_count(), torch.cuda.memory_allocated())) # ----------------------------- data done -------------------------------- # ------------------------ parepare optimizer, scheduler, criterion ------- optimizer = get_optimizer(engine, cfg, model, no_l2_keywords=no_l2_keywords, use_nesterov=use_nesterov, keyword_to_lr_mult=keyword_to_lr_mult) scheduler = get_lr_scheduler(cfg, optimizer) criterion = get_criterion(cfg).cuda() # --------------------------------- done ------------------------------- engine.register_state(scheduler=scheduler, model=model, optimizer=optimizer) if engine.distributed: torch.cuda.set_device(local_rank) engine.echo('Distributed training, device {}'.format(local_rank)) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], broadcast_buffers=False, ) else: assert torch.cuda.device_count() == 1 engine.echo('Single GPU training') if tensorflow_style_init: init_as_tensorflow(model) if cfg.init_weights: engine.load_checkpoint(cfg.init_weights) if init_hdf5: engine.load_hdf5(init_hdf5, load_weights_keyword=load_weights_keyword) if auto_continue: assert cfg.init_weights is None engine.load_checkpoint(get_last_checkpoint(cfg.output_dir)) if show_variables: engine.show_variables() # ------------ do training ---------------------------- # engine.log("\n\nStart training with pytorch version {}".format( torch.__version__)) iteration = engine.state.iteration iters_per_epoch = num_iters_per_epoch(cfg) max_iters = iters_per_epoch * cfg.max_epochs tb_writer = SummaryWriter(cfg.tb_dir) tb_tags = ['Top1-Acc', 'Top5-Acc', 'Loss'] model.train() done_epochs = iteration // iters_per_epoch last_epoch_done_iters = iteration % iters_per_epoch if done_epochs == 0 and last_epoch_done_iters == 0: engine.save_hdf5(os.path.join(cfg.output_dir, 'init.hdf5')) recorded_train_time = 0 recorded_train_examples = 0 collected_train_loss_sum = 0 collected_train_loss_count = 0 if gradient_mask is not None: gradient_mask_tensor = {} for name, value in gradient_mask.items(): gradient_mask_tensor[name] = torch.Tensor(value).cuda() else: gradient_mask_tensor = None for epoch in range(done_epochs, cfg.max_epochs): if engine.distributed and hasattr(train_data, 'train_sampler'): train_data.train_sampler.set_epoch(epoch) if epoch == done_epochs: pbar = tqdm(range(iters_per_epoch - last_epoch_done_iters)) else: pbar = tqdm(range(iters_per_epoch)) if epoch == 0 and local_rank == 0: val_during_train(epoch=epoch, iteration=iteration, tb_tags=tb_tags, engine=engine, model=model, val_data=val_data, criterion=criterion, descrip_str='Init', dataset_name=cfg.dataset_name, test_batch_size=TEST_BATCH_SIZE, tb_writer=tb_writer) top1 = AvgMeter() top5 = AvgMeter() losses = AvgMeter() discrip_str = 'Epoch-{}/{}'.format(epoch, cfg.max_epochs) pbar.set_description('Train' + discrip_str) for _ in pbar: start_time = time.time() data, label = load_cuda_data(train_data, dataset_name=cfg.dataset_name) # load_cuda_data(train_dataloader, cfg.dataset_name) data_time = time.time() - start_time if_accum_grad = ((iteration % cfg.grad_accum_iters) != 0) train_net_time_start = time.time() acc, acc5, loss = train_one_step( model, data, label, optimizer, criterion, if_accum_grad, gradient_mask_tensor=gradient_mask_tensor, lasso_keyword_to_strength=lasso_keyword_to_strength) train_net_time_end = time.time() if iteration > TRAIN_SPEED_START * max_iters and iteration < TRAIN_SPEED_END * max_iters: recorded_train_examples += cfg.global_batch_size recorded_train_time += train_net_time_end - train_net_time_start scheduler.step() for module in model.modules(): if hasattr(module, 'set_cur_iter'): module.set_cur_iter(iteration) if iteration % cfg.tb_iter_period == 0 and engine.world_rank == 0: for tag, value in zip( tb_tags, [acc.item(), acc5.item(), loss.item()]): tb_writer.add_scalars(tag, {'Train': value}, iteration) top1.update(acc.item()) top5.update(acc5.item()) losses.update(loss.item()) if epoch >= cfg.max_epochs - COLLECT_TRAIN_LOSS_EPOCHS: collected_train_loss_sum += loss.item() collected_train_loss_count += 1 pbar_dic = OrderedDict() pbar_dic['data-time'] = '{:.2f}'.format(data_time) pbar_dic['cur_iter'] = iteration pbar_dic['lr'] = scheduler.get_lr()[0] pbar_dic['top1'] = '{:.5f}'.format(top1.mean) pbar_dic['top5'] = '{:.5f}'.format(top5.mean) pbar_dic['loss'] = '{:.5f}'.format(losses.mean) pbar.set_postfix(pbar_dic) iteration += 1 if iteration >= max_iters or iteration % cfg.ckpt_iter_period == 0: engine.update_iteration(iteration) if (not engine.distributed) or (engine.distributed and engine.world_rank == 0): engine.save_and_link_checkpoint(cfg.output_dir) if iteration >= max_iters: break # do something after an epoch? engine.update_iteration(iteration) engine.save_latest_ckpt(cfg.output_dir) if (epoch + 1) % save_hdf5_epochs == 0: engine.save_hdf5( os.path.join(cfg.output_dir, 'epoch-{}.hdf5'.format(epoch))) if local_rank == 0 and \ cfg.val_epoch_period > 0 and (epoch >= cfg.max_epochs - 10 or epoch % cfg.val_epoch_period == 0): val_during_train(epoch=epoch, iteration=iteration, tb_tags=tb_tags, engine=engine, model=model, val_data=val_data, criterion=criterion, descrip_str=discrip_str, dataset_name=cfg.dataset_name, test_batch_size=TEST_BATCH_SIZE, tb_writer=tb_writer) if iteration >= max_iters: break # do something after the training if recorded_train_time > 0: exp_per_sec = recorded_train_examples / recorded_train_time else: exp_per_sec = 0 engine.log( 'TRAIN speed: from {} to {} iterations, batch_size={}, examples={}, total_net_time={:.4f}, examples/sec={}' .format(int(TRAIN_SPEED_START * max_iters), int(TRAIN_SPEED_END * max_iters), cfg.global_batch_size, recorded_train_examples, recorded_train_time, exp_per_sec)) if cfg.save_weights: engine.save_checkpoint(cfg.save_weights) print('NOTE: training finished, saved to {}'.format( cfg.save_weights)) engine.save_hdf5(os.path.join(cfg.output_dir, 'finish.hdf5')) if collected_train_loss_count > 0: engine.log( 'TRAIN LOSS collected over last {} epochs: {:.6f}'.format( COLLECT_TRAIN_LOSS_EPOCHS, collected_train_loss_sum / collected_train_loss_count))
def aofp_train_main(local_rank, target_layers, succ_strategy, warmup_iterations, aofp_batches_per_half, flops_func, cfg: BaseConfigByEpoch, net=None, train_dataloader=None, val_dataloader=None, show_variables=False, convbuilder=None, init_hdf5=None, no_l2_keywords='depth', gradient_mask=None, use_nesterov=False, tensorflow_style_init=False, keyword_to_lr_mult=None, auto_continue=False, lasso_keyword_to_strength=None, save_hdf5_epochs=10000, remain_flops_ratio=0): if no_l2_keywords is None: no_l2_keywords = [] if type(no_l2_keywords) is not list: no_l2_keywords = [no_l2_keywords] ensure_dir(cfg.output_dir) ensure_dir(cfg.tb_dir) with Engine(local_rank=local_rank) as engine: engine.setup_log(name='train', log_dir=cfg.output_dir, file_name='log.txt') # ----------------------------- build model ------------------------------ if convbuilder is None: convbuilder = ConvBuilder(base_config=cfg) if net is None: net_fn = get_model_fn(cfg.dataset_name, cfg.network_type) model = net_fn(cfg, convbuilder) else: model = net model = model.cuda() # ----------------------------- model done ------------------------------ # ---------------------------- prepare data ------------------------- if train_dataloader is None: train_data = create_dataset(cfg.dataset_name, cfg.dataset_subset, cfg.global_batch_size, distributed=engine.distributed) if cfg.val_epoch_period > 0 and val_dataloader is None: val_data = create_dataset(cfg.dataset_name, 'val', global_batch_size=100, distributed=False) engine.echo('NOTE: Data prepared') engine.echo( 'NOTE: We have global_batch_size={} on {} GPUs, the allocated GPU memory is {}' .format(cfg.global_batch_size, torch.cuda.device_count(), torch.cuda.memory_allocated())) # ----------------------------- data done -------------------------------- # ------------------------ parepare optimizer, scheduler, criterion ------- optimizer = get_optimizer(engine, cfg, model, no_l2_keywords=no_l2_keywords, use_nesterov=use_nesterov, keyword_to_lr_mult=keyword_to_lr_mult) scheduler = get_lr_scheduler(cfg, optimizer) criterion = get_criterion(cfg).cuda() # --------------------------------- done ------------------------------- engine.register_state(scheduler=scheduler, model=model, optimizer=optimizer) if engine.distributed: torch.cuda.set_device(local_rank) engine.echo('Distributed training, device {}'.format(local_rank)) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], broadcast_buffers=False, ) else: assert torch.cuda.device_count() == 1 engine.echo('Single GPU training') if tensorflow_style_init: init_as_tensorflow(model) if cfg.init_weights: engine.load_checkpoint(cfg.init_weights) if init_hdf5: engine.load_part('base_path.', init_hdf5) if auto_continue: assert cfg.init_weights is None engine.load_checkpoint(get_last_checkpoint(cfg.output_dir)) if show_variables: engine.show_variables() # ------------ do training ---------------------------- # engine.log("\n\nStart training with pytorch version {}".format( torch.__version__)) iteration = engine.state.iteration iters_per_epoch = num_iters_per_epoch(cfg) max_iters = iters_per_epoch * cfg.max_epochs tb_writer = SummaryWriter(cfg.tb_dir) tb_tags = ['Top1-Acc', 'Top5-Acc', 'Loss'] model.train() done_epochs = iteration // iters_per_epoch last_epoch_done_iters = iteration % iters_per_epoch if done_epochs == 0 and last_epoch_done_iters == 0: engine.save_hdf5(os.path.join(cfg.output_dir, 'init.hdf5')) recorded_train_time = 0 recorded_train_examples = 0 collected_train_loss_sum = 0 collected_train_loss_count = 0 if gradient_mask is not None: gradient_mask_tensor = {} for name, value in gradient_mask.items(): gradient_mask_tensor[name] = torch.Tensor(value).cuda() else: gradient_mask_tensor = None ######################### aofp _init_interval = aofp_batches_per_half // len(target_layers) layer_to_start_iter = { i: (_init_interval * i + warmup_iterations) for i in target_layers } print( 'the initial layer_to_start_iter = {}'.format(layer_to_start_iter)) # 0. get all the AOFPLayers layer_idx_to_module = {} for submodule in model.modules(): if hasattr(submodule, 'score_mask') or hasattr( submodule, 't_value'): layer_idx_to_module[submodule.conv_idx] = submodule print(layer_idx_to_module) ###################################### for epoch in range(done_epochs, cfg.max_epochs): if engine.distributed and hasattr(train_data, 'train_sampler'): train_data.train_sampler.set_epoch(epoch) if epoch == done_epochs: pbar = tqdm(range(iters_per_epoch - last_epoch_done_iters)) else: pbar = tqdm(range(iters_per_epoch)) if epoch == 0 and local_rank == 0: val_during_train(epoch=epoch, iteration=iteration, tb_tags=tb_tags, engine=engine, model=model, val_data=val_data, criterion=criterion, descrip_str='Init', dataset_name=cfg.dataset_name, test_batch_size=TEST_BATCH_SIZE, tb_writer=tb_writer) top1 = AvgMeter() top5 = AvgMeter() losses = AvgMeter() discrip_str = 'Epoch-{}/{}'.format(epoch, cfg.max_epochs) pbar.set_description('Train' + discrip_str) for _ in pbar: start_time = time.time() data, label = load_cuda_data(train_data, dataset_name=cfg.dataset_name) # load_cuda_data(train_dataloader, cfg.dataset_name) data_time = time.time() - start_time if_accum_grad = ((iteration % cfg.grad_accum_iters) != 0) train_net_time_start = time.time() ############ aofp # 1. see if it is time to start on every layer # 2. forward and accumulate # 3. if a half on some layer is finished, do something # ---- fetch its accumulated t vectors, analyze the first 'granu' elements # ---- if good enough, set the base mask, reset the search space # ---- elif granu == 1, do nothing # ---- else, granu /= 2, reset the search space for layer_idx, start_iter in layer_to_start_iter.items(): if start_iter == iteration: layer_idx_to_module[layer_idx].start_aofp(iteration) acc, acc5, loss = train_one_step( model, data, label, optimizer, criterion, if_accum_grad, gradient_mask_tensor=gradient_mask_tensor, lasso_keyword_to_strength=lasso_keyword_to_strength) for layer_idx, aofp_layer in layer_idx_to_module.items(): # accumulate if layer_idx not in succ_strategy: continue follow_layer_idx = succ_strategy[layer_idx] if follow_layer_idx not in layer_idx_to_module: continue t_value = layer_idx_to_module[follow_layer_idx].t_value aofp_layer.accumulate_t_value(t_value) if aofp_layer.finished_a_half(iteration): aofp_layer.halve_or_stop(iteration) ################################### train_net_time_end = time.time() if iteration > TRAIN_SPEED_START * max_iters and iteration < TRAIN_SPEED_END * max_iters: recorded_train_examples += cfg.global_batch_size recorded_train_time += train_net_time_end - train_net_time_start scheduler.step() for module in model.modules(): if hasattr(module, 'set_cur_iter'): module.set_cur_iter(iteration) if iteration % cfg.tb_iter_period == 0 and engine.world_rank == 0: for tag, value in zip( tb_tags, [acc.item(), acc5.item(), loss.item()]): tb_writer.add_scalars(tag, {'Train': value}, iteration) top1.update(acc.item()) top5.update(acc5.item()) losses.update(loss.item()) if epoch >= cfg.max_epochs - COLLECT_TRAIN_LOSS_EPOCHS: collected_train_loss_sum += loss.item() collected_train_loss_count += 1 pbar_dic = OrderedDict() pbar_dic['data-time'] = '{:.2f}'.format(data_time) pbar_dic['cur_iter'] = iteration pbar_dic['lr'] = scheduler.get_lr()[0] pbar_dic['top1'] = '{:.5f}'.format(top1.mean) pbar_dic['top5'] = '{:.5f}'.format(top5.mean) pbar_dic['loss'] = '{:.5f}'.format(losses.mean) pbar.set_postfix(pbar_dic) iteration += 1 if iteration >= max_iters or iteration % cfg.ckpt_iter_period == 0: engine.update_iteration(iteration) if (not engine.distributed) or (engine.distributed and engine.world_rank == 0): engine.save_and_link_checkpoint(cfg.output_dir) if iteration >= max_iters: break # do something after an epoch? engine.update_iteration(iteration) engine.save_latest_ckpt(cfg.output_dir) if (epoch + 1) % save_hdf5_epochs == 0: engine.save_hdf5( os.path.join(cfg.output_dir, 'epoch-{}.hdf5'.format(epoch))) if local_rank == 0 and \ cfg.val_epoch_period > 0 and (epoch >= cfg.max_epochs - 10 or epoch % cfg.val_epoch_period == 0): val_during_train(epoch=epoch, iteration=iteration, tb_tags=tb_tags, engine=engine, model=model, val_data=val_data, criterion=criterion, descrip_str=discrip_str, dataset_name=cfg.dataset_name, test_batch_size=TEST_BATCH_SIZE, tb_writer=tb_writer) cur_deps = np.array(cfg.deps) for submodule in model.modules(): if hasattr(submodule, 'base_mask'): cur_deps[submodule.conv_idx] = np.sum( submodule.base_mask.cpu().numpy() == 1) origin_flops = flops_func(cfg.deps) cur_flops = flops_func(cur_deps) remain_ratio = cur_flops / origin_flops if local_rank == 0: print('##########################') print('origin deps ', cfg.deps) print('cur deps ', cur_deps) print('remain flops ratio = ', remain_ratio, 'the target is ', remain_flops_ratio) print('##########################') if remain_ratio < remain_flops_ratio: break if iteration >= max_iters: break # do something after the training if recorded_train_time > 0: exp_per_sec = recorded_train_examples / recorded_train_time else: exp_per_sec = 0 engine.log( 'TRAIN speed: from {} to {} iterations, batch_size={}, examples={}, total_net_time={:.4f}, examples/sec={}' .format(int(TRAIN_SPEED_START * max_iters), int(TRAIN_SPEED_END * max_iters), cfg.global_batch_size, recorded_train_examples, recorded_train_time, exp_per_sec)) if cfg.save_weights: engine.save_checkpoint(cfg.save_weights) print('NOTE: training finished, saved to {}'.format( cfg.save_weights)) engine.save_hdf5(os.path.join(cfg.output_dir, 'finish.hdf5')) if collected_train_loss_count > 0: engine.log( 'TRAIN LOSS collected over last {} epochs: {:.6f}'.format( COLLECT_TRAIN_LOSS_EPOCHS, collected_train_loss_sum / collected_train_loss_count)) final_deps = aofp_prune(model, origin_deps=cfg.deps, succ_strategy=succ_strategy, save_path=os.path.join(cfg.output_dir, 'finish_pruned.hdf5')) origin_flops = flops_func(cfg.deps) cur_flops = flops_func(final_deps) engine.log( '##################################################################' ) engine.log(cfg.network_type) engine.log('origin width: {} , flops {} '.format( cfg.deps, origin_flops)) engine.log('final width: {}, flops {} '.format(final_deps, cur_flops)) engine.log('flops reduction: {}'.format(1 - cur_flops / origin_flops)) return final_deps
def train(cfg): """ Train a video model for many epochs on train set and evaluate it on val set. Args: cfg (CfgNode): configs. Details can be found in slowfast/config/defaults.py """ # Set up environment. du.init_distributed_training(cfg) # Set random seed from configs. np.random.seed(cfg.RNG_SEED) torch.manual_seed(cfg.RNG_SEED) # Setup logging format. logging.setup_logging(cfg.OUTPUT_DIR) # Init multigrid. multigrid = None if cfg.MULTIGRID.LONG_CYCLE or cfg.MULTIGRID.SHORT_CYCLE: multigrid = MultigridSchedule() cfg = multigrid.init_multigrid(cfg) if cfg.MULTIGRID.LONG_CYCLE: cfg, _ = multigrid.update_long_cycle(cfg, cur_epoch=0) # Print config. logger.info("Train with config:") logger.info(pprint.pformat(cfg)) # Build the video model and print model statistics. model = build_model(cfg) # model = x3d.MyModel() if du.is_master_proc() and cfg.LOG_MODEL_INFO: misc.log_model_info(model, cfg, is_train=True) # Construct the optimizer. optimizer = optim.construct_optimizer(model, cfg) # Load a checkpoint to resume training if applicable. if cfg.TRAIN.AUTO_RESUME and cu.has_checkpoint(cfg.OUTPUT_DIR): last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR) logger.info("Load from last checkpoint, {}.".format(last_checkpoint)) checkpoint_epoch = cu.load_checkpoint( last_checkpoint, model, cfg.NUM_GPUS > 1, optimizer ) start_epoch = checkpoint_epoch + 1 elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "": logger.info("Load from given checkpoint file.") checkpoint_epoch = cu.load_checkpoint( cfg.TRAIN.CHECKPOINT_FILE_PATH, model, cfg.NUM_GPUS > 1, optimizer, inflation=cfg.TRAIN.CHECKPOINT_INFLATE, convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2", ) start_epoch = checkpoint_epoch + 1 else: start_epoch = 0 # Create the video train and val loaders. train_loader = loader.construct_loader(cfg, "train") val_loader = loader.construct_loader(cfg, "val") precise_bn_loader = loader.construct_loader( cfg, "train", is_precise_bn=True ) # Create meters. if cfg.DETECTION.ENABLE: train_meter = AVAMeter(len(train_loader), cfg, mode="train") val_meter = AVAMeter(len(val_loader), cfg, mode="val") else: train_meter = TrainMeter(len(train_loader), cfg) val_meter = ValMeter(len(val_loader), cfg) # set up writer for logging to Tensorboard format. if cfg.TENSORBOARD.ENABLE and du.is_master_proc( cfg.NUM_GPUS * cfg.NUM_SHARDS ): writer = tb.TensorboardWriter(cfg) else: writer = None # Perform the training loop. logger.info("Start epoch: {}".format(start_epoch + 1)) for cur_epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCH): if cfg.MULTIGRID.LONG_CYCLE: cfg, changed = multigrid.update_long_cycle(cfg, cur_epoch) if changed: ( model, optimizer, train_loader, val_loader, precise_bn_loader, train_meter, val_meter, ) = build_trainer(cfg) # Load checkpoint. if cu.has_checkpoint(cfg.OUTPUT_DIR): last_checkpoint = cu.get_last_checkpoint(cfg.OUTPUT_DIR) assert "{:05d}.pyth".format(cur_epoch) in last_checkpoint else: last_checkpoint = cfg.TRAIN.CHECKPOINT_FILE_PATH logger.info("Load from {}".format(last_checkpoint)) cu.load_checkpoint( last_checkpoint, model, cfg.NUM_GPUS > 1, optimizer ) # Shuffle the dataset. loader.shuffle_dataset(train_loader, cur_epoch) # Train for one epoch. train_epoch( train_loader, model, optimizer, train_meter, cur_epoch, cfg, writer ) # Compute precise BN stats. if cfg.BN.USE_PRECISE_STATS and len(get_bn_modules(model)) > 0: calculate_and_update_precise_bn( precise_bn_loader, model, min(cfg.BN.NUM_BATCHES_PRECISE, len(precise_bn_loader)), ) _ = misc.aggregate_sub_bn_stats(model) # Save a checkpoint. if cu.is_checkpoint_epoch( cfg, cur_epoch, None if multigrid is None else multigrid.schedule ): cu.save_checkpoint(cfg.OUTPUT_DIR, model, optimizer, cur_epoch, cfg) # Evaluate the model on validation set. if misc.is_eval_epoch( cfg, cur_epoch, None if multigrid is None else multigrid.schedule ): eval_epoch(val_loader, model, val_meter, cur_epoch, cfg, writer) if writer is not None: writer.close()