def analyze_single_model(config_file_path, args, plot=True): cpu_device = torch.device('cpu') if file_util.check_if_exists(config_file_path): config = yaml_util.load_yaml_file(config_file_path) if 'teacher_model' in config: teacher_model_config = config['teacher_model'] org_model, teacher_model_type = mimic_util.get_org_model( teacher_model_config, cpu_device) model = mimic_util.get_mimic_model(config, org_model, teacher_model_type, teacher_model_config, cpu_device) model_type = config['mimic_model']['type'] input_shape = config['input_shape'] elif 'autoencoder' in config: model, model_type = ae_util.get_autoencoder( config, cpu_device, True) input_shape = config['input_shape'] else: model, model_type, input_shape = read_config(config_file_path) else: pickle_file_path = args.pkl model_type = args.model input_shape = list(data_util.convert2type_list(args.isize, ',', int)) model = file_util.load_pickle(pickle_file_path) if file_util.check_if_exists(pickle_file_path)\ else get_model(model_type) op_counts, data_sizes, accum_complexities =\ analyze(model, input_shape, model_type, args.scale, args.submodule, plot) return op_counts, data_sizes, accum_complexities, model_type
def run(args): device = torch.device('cuda' if torch.cuda.is_available() and not args.cpu else 'cpu') if device.type == 'cuda': cudnn.benchmark = True config = yaml_util.load_yaml_file(args.config) dataset_config = config['dataset'] train_config = config['train'] test_config = config['test'] compress_config = test_config['compression'] input_shape = config['input_shape'] train_loader, valid_loader, test_loader =\ dataset_util.get_data_loaders(dataset_config, batch_size=train_config['batch_size'], compression_type=compress_config['type'], compressed_size=compress_config['size'], rough_size=train_config['rough_size'], reshape_size=input_shape[1:3], test_batch_size=test_config['batch_size'], jpeg_quality=test_config['jquality']) pickle_file_path = args.pkl if not file_util.check_if_exists(pickle_file_path): model = module_util.get_model(config, device) resume_from_ckpt(model, config['model'], device) else: model = file_util.load_pickle(pickle_file_path).to(device) analysis_mode = args.mode model.eval() if analysis_mode == 'comp_rate': analyze_compression_rate(model, input_shape, test_loader, device) elif analysis_mode == 'run_time': analyze_running_time(model, input_shape, args.comp_layer, test_loader, device) else: raise ValueError('mode argument `{}` is not expected'.format(analysis_mode))
def resume_from_ckpt(ckpt_file_path, model, is_student=False): if not file_util.check_if_exists(ckpt_file_path): print('{} checkpoint was not found at {}'.format( "Student" if is_student else "Teacher", ckpt_file_path)) if is_student: return 1, None return 1 print('Resuming from checkpoint..') ckpt = torch.load(ckpt_file_path) state_dict = ckpt['model'] if not is_student and isinstance(model, Inception3) or\ (hasattr(model, 'module') and isinstance(model.module, Inception3)): for key in list(state_dict.keys()): if key.startswith('AuxLogits') or key.startswith( 'module.AuxLogits'): state_dict.pop(key) model.load_state_dict(state_dict) start_epoch = ckpt['epoch'] if is_student: return start_epoch, ckpt[ 'best_avg_loss'] if 'best_avg_loss' in ckpt else ckpt[ 'best_valid_value'] return start_epoch
def train(model, train_sampler, train_data_loader, val_data_loader, device, distributed, config, args, ckpt_file_path): train_config = config['train'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) best_val_map = 0.0 if file_util.check_if_exists(ckpt_file_path): best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) num_epochs = train_config['num_epochs'] log_freq = train_config['log_freq'] start_time = time.time() for epoch in range(num_epochs): if distributed: train_sampler.set_epoch(epoch) train_model(model, optimizer, train_data_loader, device, epoch, log_freq) lr_scheduler.step() # evaluate after every epoch coco_evaluator = main_util.evaluate(model, val_data_loader, device=device) # Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] val_map = coco_evaluator.coco_eval['bbox'].stats[0] if val_map > best_val_map: print('Updating ckpt (Best BBox mAP: {:.4f} -> {:.4f})'.format(best_val_map, val_map)) best_val_map = val_map save_ckpt(model, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path) lr_scheduler.step() dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, device, distributed, distill_backbone_only, config, args): train_config = config['train'] distillation_box = DistillationBox(teacher_model, student_model, train_config['criterion']) ckpt_file_path = config['student_model']['ckpt'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) use_bottleneck_transformer = args.transform_bottleneck best_val_map = 0.0 if file_util.check_if_exists(ckpt_file_path): best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) num_epochs = train_config['num_epochs'] log_freq = train_config['log_freq'] teacher_model_without_dp = teacher_model.module if isinstance( teacher_model, DataParallel) else teacher_model student_model_without_ddp =\ student_model.module if isinstance(student_model, DistributedDataParallel) else student_model start_time = time.time() for epoch in range(num_epochs): if distributed: train_sampler.set_epoch(epoch) teacher_model.eval() student_model.train() teacher_model_without_dp.distill_backbone_only = distill_backbone_only student_model_without_ddp.distill_backbone_only = distill_backbone_only student_model_without_ddp.backbone.body.layer1.use_bottleneck_transformer = False distill_model(distillation_box, train_data_loader, optimizer, log_freq, device, epoch) student_model_without_ddp.distill_backbone_only = False student_model_without_ddp.backbone.body.layer1.use_bottleneck_transformer = use_bottleneck_transformer coco_evaluator = main_util.evaluate(student_model, val_data_loader, device=device) # Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] val_map = coco_evaluator.coco_eval['bbox'].stats[0] if val_map > best_val_map and misc_util.is_main_process(): print('Updating ckpt (Best BBox mAP: {:.4f} -> {:.4f})'.format( best_val_map, val_map)) best_val_map = val_map save_ckpt(student_model_without_ddp, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path) lr_scheduler.step() dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def resume_from_ckpt(ckpt_file_path, autoencoder): if not file_util.check_if_exists(ckpt_file_path): print('Autoencoder checkpoint was not found at {}'.format( ckpt_file_path)) return 1, None print('Resuming from checkpoint..') checkpoint = torch.load(ckpt_file_path) state_dict = checkpoint['model'] autoencoder.load_state_dict(state_dict) start_epoch = checkpoint['epoch'] return start_epoch, checkpoint['best_value']
def resume_from_ckpt(model, model_config, init): ckpt_file_path = model_config['ckpt'] if init or not file_util.check_if_exists(ckpt_file_path): return model_config['type'], 0, 1, ckpt_file_path print('Resuming from checkpoint..') checkpoint = torch.load(ckpt_file_path) model.load_state_dict(checkpoint['model']) model_type = checkpoint['type'] best_acc = checkpoint['acc'] start_epoch = checkpoint['epoch'] return model_type, best_acc, start_epoch, ckpt_file_path
def distill(teacher_model, student_model, train_data_loader, val_data_loader, device, distributed, start_epoch, config, args): print('Start knowledge distillation') train_config = config['train'] distillation_box = DistillationBox(teacher_model, student_model, train_config['criterion']) ckpt_file_path = config['mimic_model']['ckpt'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) best_val_top1_accuracy = 0.0 if file_util.check_if_exists(ckpt_file_path): best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) interval = train_config['interval'] if interval <= 0: num_batches = len(train_data_loader) interval = num_batches // 20 if num_batches >= 20 else 1 student_model_without_ddp = \ student_model.module if isinstance(student_model, DistributedDataParallel) else student_model start_time = time.time() for epoch in range(start_epoch, train_config['epoch']): if distributed: train_data_loader.sampler.set_epoch(epoch) teacher_model.eval() student_model.train() distill_one_epoch(distillation_box, train_data_loader, optimizer, device, epoch, interval, args.apex) val_top1_accuracy =\ evaluate(student_model, val_data_loader, device=device, interval=interval, split_name='Validation') if val_top1_accuracy > best_val_top1_accuracy and main_util.is_main_process( ): print( 'Updating ckpt (Best top1 accuracy: {:.4f} -> {:.4f})'.format( best_val_top1_accuracy, val_top1_accuracy)) best_val_top1_accuracy = val_top1_accuracy save_ckpt(student_model_without_ddp, optimizer, lr_scheduler, best_val_top1_accuracy, config, args, ckpt_file_path) lr_scheduler.step() dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def visualize_predictions(model, input_file_paths, device, output_dir_path): model.eval() file_util.make_dirs(output_dir_path) for input_file_path in input_file_paths: if not file_util.check_if_exists(input_file_path): print('`{}` is not found.'.format(input_file_path)) continue if os.path.isfile(input_file_path): predict(model, input_file_path, device, output_dir_path) else: for sub_input_file_path in file_util.get_file_path_list( input_file_path, is_recursive=True): predict(model, sub_input_file_path, device, output_dir_path)
def distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, device, distributed, distill_backbone_only, config, args): train_config = config['train'] student_config = config['student_model'] distillation_box = DistillationBox(teacher_model, student_model, train_config['criterion'], student_config) ckpt_file_path = config['student_model']['ckpt'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) if file_util.check_if_exists(ckpt_file_path): best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) save_ckpt(student_model, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path)
def load_ckpt(ckpt_file_path, model=None, optimizer=None, lr_scheduler=None, strict=True, require_weights=False): if not file_util.check_if_exists(ckpt_file_path): print('ckpt file is not found at `{}`'.format(ckpt_file_path)) if require_weights: raise ValueError('Ckpt file is not found at `{}`'.format(ckpt_file_path)) return None, None ckpt = torch.load(ckpt_file_path, map_location='cpu') if model is not None: print('Loading model parameters') model.load_state_dict(ckpt['model'], strict=strict) if optimizer is not None: print('Loading optimizer parameters') optimizer.load_state_dict(ckpt['optimizer']) if lr_scheduler is not None: print('Loading scheduler parameters') lr_scheduler.load_state_dict(ckpt['lr_scheduler']) return ckpt.get('best_value', 0.0), ckpt['config'], ckpt['args']
def resume_from_ckpt(model, model_config, device): ckpt_file_path = model_config['ckpt'] if not file_util.check_if_exists(ckpt_file_path): return model_config['type'], 0, 1, ckpt_file_path print('Resuming from checkpoint..') checkpoint = torch.load(ckpt_file_path) model_state_dict = checkpoint['model'] if device.type == 'cpu': for key in list(model_state_dict.keys()): if key.startswith('module.'): val = model_state_dict.pop(key) model_state_dict[key[7:]] = val model.load_state_dict(model_state_dict) model_type = checkpoint['type'] best_acc = checkpoint['acc'] start_epoch = checkpoint['epoch'] return model_type, best_acc, start_epoch, ckpt_file_path
def train(model, ext_classifier, train_sampler, train_data_loader, val_data_loader, device, distributed, config, args, ckpt_file_path): train_config = config['train'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(ext_classifier, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) best_val_roc_auc = 0.0 if file_util.check_if_exists(ckpt_file_path): best_val_roc_auc, _, _ =\ load_ckpt(ckpt_file_path, model=ext_classifier, optimizer=optimizer, lr_scheduler=lr_scheduler) num_epochs = train_config['num_epochs'] log_freq = train_config['log_freq'] start_time = time.time() for epoch in range(num_epochs): if distributed: train_sampler.set_epoch(epoch) train_model(model, optimizer, train_data_loader, device, epoch, log_freq) lr_scheduler.step() # evaluate after every epoch val_roc_auc = evaluate(model, val_data_loader, device, min_recall=args.min_recall, split_name='Validation') if val_roc_auc > best_val_roc_auc: print('Updating ckpt (Best ROC-AUC: {:.4f} -> {:.4f})'.format( best_val_roc_auc, val_roc_auc)) best_val_roc_auc = val_roc_auc save_ckpt(ext_classifier, optimizer, lr_scheduler, best_val_roc_auc, config, args, ckpt_file_path) dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
def main(args): config = yaml_util.load_yaml_file(args.config) if args.json is not None: main_util.overwrite_config(config, args.json) distributed, device_ids = main_util.init_distributed_mode(args.world_size, args.dist_url) device = torch.device(args.device if torch.cuda.is_available() else 'cpu') teacher_model = get_model(config['teacher_model'], device) module_util.freeze_module_params(teacher_model) student_model_config = config['student_model'] student_model = get_model(student_model_config, device) freeze_modules(student_model, student_model_config) ckpt_file_path = config['student_model']['ckpt'] train_config = config['train'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) if file_util.check_if_exists(ckpt_file_path): best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) save_ckpt(student_model, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path)
def distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, device, distributed, distill_backbone_only, config, args): train_config = config['train'] student_config = config['student_model'] distillation_box = DistillationBox(teacher_model, student_model, train_config['criterion'], student_config) ckpt_file_path = config['student_model']['ckpt'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) scheduler_config = train_config['scheduler'] lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params']) use_bottleneck_transformer = args.transform_bottleneck best_val_map = 0.0 if file_util.check_if_exists(ckpt_file_path): if args.ignore_optimizer: best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=None, lr_scheduler=None) else: best_val_map, _, _ = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler) num_epochs = train_config['num_epochs'] log_freq = train_config['log_freq'] teacher_model_without_dp = teacher_model.module if isinstance( teacher_model, DataParallel) else teacher_model student_model_without_ddp = \ student_model.module if isinstance(student_model, DistributedDataParallel) else student_model start_time = time.time() post_bn = False if 'post_batch_norm' in config['train']: post_bn = config['train']['post_batch_norm'] for epoch in range(lr_scheduler.last_epoch, num_epochs): if distributed: train_sampler.set_epoch(epoch) teacher_model.eval() student_model.train() teacher_model_without_dp.distill_backbone_only = distill_backbone_only student_model_without_ddp.distill_backbone_only = distill_backbone_only set_bottleneck_transformer(student_model_without_ddp, False) distill_model(distillation_box, train_data_loader, optimizer, log_freq, device, epoch) student_model_without_ddp.distill_backbone_only = False set_bottleneck_transformer(student_model_without_ddp, use_bottleneck_transformer) val_map = 0 width_list = [1.0] if 'slimmable' in student_config['backbone']['params']: width_list = [0.25, 0.5, 0.75, 1.0] width_list = [ w for w in width_list if w in student_config['backbone']['params']['width_mult_list'] ] for width in width_list: set_width(student_model, width) if post_bn: ComputeBN(student_model, train_data_loader) print('\n[Student model@width={}]'.format(width)) coco_evaluator = main_util.evaluate(student_model, val_data_loader, device=device) val_map += coco_evaluator.coco_eval['bbox'].stats[0] val_map = val_map / len(width_list) print('BBox mAP: {:.4f})'.format(val_map)) if val_map > best_val_map and misc_util.is_main_process(): print('Updating ckpt (Best BBox mAP: {:.4f} -> {:.4f})'.format( best_val_map, val_map)) best_val_map = val_map save_ckpt(student_model_without_ddp, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path) lr_scheduler.step() if distributed: dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))