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 run(args): 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') if device.type == 'cuda': cudnn.benchmark = True print(args) config = yaml_util.load_yaml_file(args.config) input_shape = config['input_shape'] ckpt_file_path = config['autoencoder']['ckpt'] train_loader, valid_loader, test_loader = main_util.get_data_loaders( config, distributed) if not args.test_only: train(train_loader, valid_loader, input_shape, config, device, distributed, device_ids) autoencoder, _ = ae_util.get_autoencoder(config, device) resume_from_ckpt(ckpt_file_path, autoencoder) extended_model, model = ae_util.get_extended_model(autoencoder, config, input_shape, device) if not args.extended_only: if device.type == 'cuda': model = DistributedDataParallel(model, device_ids=device_ids) if distributed \ else DataParallel(model) evaluate(model, test_loader, device, title='[Original model]') if device.type == 'cuda': extended_model = DistributedDataParallel(extended_model, device_ids=device_ids) if distributed \ else DataParallel(extended_model) evaluate(extended_model, test_loader, device, title='[Mimic model]')
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 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) print('Updatable parameters: {}'.format( module_util.get_updatable_param_names(student_model))) distill_backbone_only = student_model_config['distill_backbone_only'] train_config = config['train'] train_sampler, train_data_loader, val_data_loader, test_data_loader = \ data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) if distributed: teacher_model = DataParallel(teacher_model, device_ids=device_ids) student_model = DistributedDataParallel(student_model, device_ids=device_ids) if args.distill: distill(teacher_model, student_model, train_sampler, train_data_loader, val_data_loader, device, distributed, distill_backbone_only, config, args) load_ckpt( config['student_model']['ckpt'], model=student_model.module if isinstance( student_model, DistributedDataParallel) else student_model) evaluate(teacher_model, student_model, test_data_loader, device, args.skip_teacher_eval, args.transform_bottleneck)
def main(args): distributed, device_ids = main_util.init_distributed_mode( args.world_size, args.dist_url) config = yaml_util.load_yaml_file(args.config) if args.json is not None: main_util.overwrite_config(config, args.json) device = torch.device(args.device) print(args) print('Loading data') train_config = config['train'] train_sampler, train_data_loader, val_data_loader, test_data_loader =\ data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) print('Creating model') model_config = config['model'] model = get_model(model_config, device) print('Model Created') if distributed: model = nn.parallel.DistributedDataParallel(model, device_ids=device_ids) if args.train: print('Start training') start_time = time.time() train(model, train_sampler, train_data_loader, val_data_loader, device, distributed, config, args, model_config['ckpt']) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) main_util.evaluate(model, test_data_loader, device=device)
def main(args): config = yaml_util.load_yaml_file(args.config) if args.json is not None: main_util.overwrite_config(config, args.json) device = torch.device(args.device) print(args) model_config = config.get('model', None) if args.model_params and model_config is not None: model = get_model(model_config, device) analyze_model_params(model, args.modules) if args.data_size is not None: analyze_data_size(config['dataset'], split_name=args.data_size, resized=args.resized) student_model_config = config.get('student_model', None) if args.bottleneck_size is not None and (student_model_config is not None or (model_config is not None and 'ext_config' in model_config['backbone'])): data_logger = DataLogger() tmp_model_config = student_model_config if student_model_config is not None else model_config model = get_model(tmp_model_config, device, bottleneck_transformer=data_logger) model.backbone.body.layer1.use_bottleneck_transformer = True analyze_bottleneck_size(model, data_logger, device, config['dataset'], split_name=args.bottleneck_size) if student_model_config is not None: model_config = student_model_config if args.split_model is not None: model = get_model(model_config, device) analyze_split_model_inference(model, device, args.quantize, args.skip_tail, config['dataset'], args.split_model)
def run(args): distributed, device_ids = main_util.init_distributed_mode( args.world_size, args.dist_url) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device.type == 'cuda': cudnn.benchmark = True print(args) config = yaml_util.load_yaml_file(args.config) train_loader, valid_loader, test_loader = main_util.get_data_loaders( config, distributed) if 'mimic_model' in config: model = mimic_util.get_mimic_model_easily(config, device) model_config = config['mimic_model'] else: model = module_util.get_model(config, device) model_config = config['model'] model_type, best_valid_acc, start_epoch, ckpt_file_path =\ module_util.resume_from_ckpt(model, model_config, args.init) train_config = config['train'] criterion_config = train_config['criterion'] criterion = func_util.get_loss(criterion_config['type'], criterion_config['params']) if not args.evaluate: train(model, train_loader, valid_loader, best_valid_acc, criterion, device, distributed, device_ids, train_config, args.epoch, start_epoch, args.lr, ckpt_file_path, model_type) test(model, test_loader, device)
def main(args): device = torch.device(args.device if torch.cuda.is_available() else 'cpu') result_dir = join(args.dir, "results") Path(result_dir).mkdir(parents=True, exist_ok=True) config_file_list = [f for f in listdir(args.dir) if isfile(join(args.dir, f))] config_file_list = [f for f in config_file_list if f.endswith('.yaml')] config_file_list.sort() for config_file in config_file_list: result_file = join(result_dir, config_file).split('.yaml')[0]+".csv" config_file = join(args.dir, config_file) print("#################################################################################") print("config_file: {},".format(config_file)) print("result_file: {},".format(result_file)) config = yaml_util.load_yaml_file(config_file) if args.resume and os.path.isfile(result_file): print("Experiment Done: Skipping") continue results = experiment(config, device, args) if args.dry_run: continue print(results) results.to_csv(result_file, float_format='%.4f')
def read_config(config_file_path): config = yaml_util.load_yaml_file(config_file_path) if config['model']['type'] == 'inception_v3': config['model']['params']['aux_logits'] = False model = module_util.get_model(config, torch.device('cpu')) model_type = config['model']['type'] input_shape = config['input_shape'] return model, model_type, input_shape
def get_org_model(teacher_model_config, device): teacher_config = yaml_util.load_yaml_file(teacher_model_config['config']) if teacher_config['model']['type'] == 'inception_v3': teacher_config['model']['params']['aux_logits'] = False model = module_util.get_model(teacher_config, device) model_config = teacher_config['model'] resume_from_ckpt(model_config['ckpt'], model) return model, model_config['type']
def get_teacher_model(teacher_model_config, input_shape, device): teacher_config = yaml_util.load_yaml_file(teacher_model_config['config']) model_config = teacher_config['model'] if model_config['type'] == 'inception_v3': model_config['params']['aux_logits'] = False model = module_util.get_model(teacher_config, device) resume_from_ckpt(model_config['ckpt'], model) return extract_teacher_model(model, input_shape, device, teacher_model_config), model_config['type']
def get_head_model(config, input_shape, device): org_model_config = config['org_model'] model_config = yaml_util.load_yaml_file(org_model_config['config']) sub_model_config = model_config['model'] if sub_model_config['type'] == 'inception_v3': sub_model_config['params']['aux_logits'] = False model = module_util.get_model(model_config, device) module_util.resume_from_ckpt(model, sub_model_config, False) return extract_head_model(model, input_shape, device, org_model_config['partition_idx'])
def analyze_teacher_student_models(mimic_config_file_paths, args): scaled = args.scale submoduled = args.submodule model_type_list = list() teacher_complexity_list = list() student_complexity_list = list() teacher_data_size_list = list() student_data_size_list = list() for mimic_config_file_path in mimic_config_file_paths: mimic_config = yaml_util.load_yaml_file(mimic_config_file_path) input_shape = mimic_config['input_shape'] teacher_model_type, teacher_model, student_model = get_teacher_and_student_models( mimic_config, input_shape) _, teacher_data_sizes, teacher_accum_complexities = analyze( teacher_model, input_shape, None, scaled=scaled, submoduled=submoduled, plot=False) _, student_data_sizes, student_accum_complexities = analyze( student_model, input_shape, None, scaled=scaled, submoduled=submoduled, plot=False) student_model_config = mimic_config['student_model'] student_model_version = student_model_config['version'] made_bottleneck = student_model_version.endswith('b') model_type_list.append('Ver.{}'.format(student_model_version)) teacher_complexity_list.append(teacher_accum_complexities[-1]) teacher_data_size_list.append(teacher_data_sizes[-1] / teacher_data_sizes[0]) bottleneck_idx = np.argmin( student_data_sizes) if made_bottleneck else -1 if student_data_sizes[bottleneck_idx] >= student_data_sizes[ 0] or not made_bottleneck: student_complexity_list.append(student_accum_complexities[-1]) student_data_size_list.append(student_data_sizes[-1] / student_data_sizes[0]) else: student_complexity_list.append( student_accum_complexities[bottleneck_idx - 1]) student_data_size_list.append(student_data_sizes[bottleneck_idx] / student_data_sizes[0]) net_measure_util.plot_teacher_and_student_complexities( teacher_complexity_list, student_complexity_list, model_type_list) net_measure_util.plot_bottleneck_data_size_vs_complexity( teacher_data_size_list, teacher_complexity_list, student_data_size_list, student_complexity_list, model_type_list)
def get_extended_model(autoencoder, config, input_shape, device, skip_bottleneck_size=False): org_model_config = config['org_model'] model_config = yaml_util.load_yaml_file(org_model_config['config']) sub_model_config = model_config['model'] if sub_model_config['type'] == 'inception_v3': sub_model_config['params']['aux_logits'] = False model = module_util.get_model(model_config, device) module_util.resume_from_ckpt(model, sub_model_config, False) return extend_model(autoencoder, model, input_shape, device, org_model_config['partition_idx'], skip_bottleneck_size), model
def main(args): config = yaml_util.load_yaml_file(args.config) if args.json is not None: main_util.overwrite_config(config, args.json) device = torch.device(args.device) print(args) print('Creating model') model_config = config['model'] if 'model' in config else config.get( 'student_model', None) if model_config is None: raise ValueError( '`{}` should contain model or student config at root'.format( args.config)) model = get_model(model_config, device, strict=False) visualize_predictions(model, args.image, device, args.output)
def main(args): distributed, device_ids = main_util.init_distributed_mode( args.world_size, args.dist_url) config = yaml_util.load_yaml_file(args.config) if args.json is not None: main_util.overwrite_config(config, args.json) device = torch.device(args.device) print(args) print('Loading data') train_config = config['train'] train_sampler, train_data_loader, val_data_loader, test_data_loader =\ data_util.get_coco_data_loaders(config['dataset'], train_config['batch_size'], distributed) print('Creating model') model_config = config['model'] model = get_model(model_config, device, strict=False) module_util.freeze_module_params(model) ext_classifier = model.get_ext_classifier() module_util.unfreeze_module_params(ext_classifier) print('Updatable parameters: {}'.format( module_util.get_updatable_param_names(model))) model.train_ext() if distributed: model = nn.parallel.DistributedDataParallel(model, device_ids=device_ids) if args.train: print('Start training') start_time = time.time() ckpt_file_path = model_config['backbone']['ext_config']['ckpt'] train(model, ext_classifier, train_sampler, train_data_loader, val_data_loader, device, distributed, config, args, ckpt_file_path) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) load_ckpt(ckpt_file_path, model=ext_classifier) evaluate(model, test_data_loader, device=device, min_recall=args.min_recall, split_name='Test')
def run(args): 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') if torch.cuda.is_available(): cudnn.benchmark = True print(args) config = yaml_util.load_yaml_file(args.config) dataset_config = config['dataset'] input_shape = config['input_shape'] train_config = config['train'] test_config = config['test'] train_loader, valid_loader, test_loader =\ dataset_util.get_data_loaders(dataset_config, batch_size=train_config['batch_size'], rough_size=train_config['rough_size'], reshape_size=input_shape[1:3], test_batch_size=test_config['batch_size'], jpeg_quality=-1, distributed=distributed) teacher_model_config = config['teacher_model'] if not args.test_only: distill(train_loader, valid_loader, input_shape, args.aux, config, device, distributed, device_ids) org_model, teacher_model_type = mimic_util.get_org_model( teacher_model_config, device) if not args.student_only: if distributed: org_model = DataParallel(org_model, device_ids=device_ids) evaluate(org_model, test_loader, device, title='[Original model]') mimic_model = mimic_util.get_mimic_model(config, org_model, teacher_model_type, teacher_model_config, device) mimic_model_without_dp = mimic_model.module if isinstance( mimic_model, DataParallel) else mimic_model file_util.save_pickle(mimic_model_without_dp, config['mimic_model']['ckpt']) if distributed: mimic_model = DistributedDataParallel(mimic_model_without_dp, device_ids=device_ids) evaluate(mimic_model, test_loader, device, title='[Mimic model]')
def run(args): print(args) config = yaml_util.load_yaml_file(args.config) sensor_device = torch.device('cpu' if args.scpu else 'cuda') edge_device = torch.device('cpu' if args.ecpu else 'cuda') partition_idx = args.partition head_output_file_path = args.head tail_output_file_path = args.tail input_shape = config['input_shape'] if 'teacher_model' not in config: model = module_util.get_model( config, torch.device('cuda') if torch.cuda.is_available() else None) module_util.resume_from_ckpt(model, config['model'], False) else: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model, teacher_model_type =\ mimic_util.get_org_model(config['teacher_model'], device) if args.org and head_output_file_path is not None and tail_output_file_path is not None: split_original_model(model, input_shape, device, config, sensor_device, edge_device, partition_idx, head_output_file_path, tail_output_file_path, args.test, args.spbit) elif args.mimic: model = mimic_util.get_mimic_model_easily(config, sensor_device) student_model_config = config['mimic_model'] load_ckpt(student_model_config['ckpt'], model=model, strict=True) split_original_model(model, input_shape, device, config, sensor_device, edge_device, partition_idx, head_output_file_path, tail_output_file_path, args.test, args.spbit) elif head_output_file_path is not None and tail_output_file_path is not None: split_within_student_model(model, input_shape, device, config, teacher_model_type, sensor_device, edge_device, partition_idx, head_output_file_path, tail_output_file_path, args.test, args.spbit) if args.model is not None and args.device is not None: convert_model(model, torch.device(args.device), args.model)
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 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) train_loader, valid_loader, test_loader = main_util.get_data_loaders( config, False) input_shape = config['input_shape'] if 'mimic_model' in config: model = mimic_util.get_mimic_model_easily(config, device) model_config = config['mimic_model'] else: model = module_util.get_model(config, device) model_config = config['model'] model_type, _, _, _ = module_util.resume_from_ckpt(model, model_config, False) split_name = args.split data_loader = train_loader if split_name == 'train' else valid_loader if split_name == 'valid' else test_loader analyze_with_mean_inputs(model, input_shape, data_loader, device, split_name, args.method, args.dim, model_type, args.output)
def main(args): if args.apex: if sys.version_info < (3, 0): raise RuntimeError( 'Apex currently only supports Python 3. Aborting.') if amp is None: raise RuntimeError( 'Failed to import apex. Please install apex from https://www.github.com/nvidia/apex ' 'to enable mixed-precision training.') distributed, device_ids = main_util.init_distributed_mode( args.world_size, args.dist_url) print(args) if torch.cuda.is_available(): torch.backends.cudnn.benchmark = True config = yaml_util.load_yaml_file(args.config) device = torch.device(args.device if torch.cuda.is_available() else 'cpu') dataset_config = config['dataset'] input_shape = config['input_shape'] train_config = config['train'] test_config = config['test'] train_data_loader, val_data_loader, test_data_loader =\ dataset_util.get_data_loaders(dataset_config, batch_size=train_config['batch_size'], rough_size=train_config['rough_size'], reshape_size=input_shape[1:3], jpeg_quality=-1, test_batch_size=test_config['batch_size'], distributed=distributed) teacher_model_config = config['teacher_model'] teacher_model, teacher_model_type = mimic_util.get_org_model( teacher_model_config, device) module_util.freeze_module_params(teacher_model) student_model = mimic_util.get_mimic_model_easily(config, device) student_model_config = config['mimic_model'] optim_config = train_config['optimizer'] optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params']) use_apex = args.apex if use_apex: student_model, optimizer = amp.initialize( student_model, optimizer, opt_level=args.apex_opt_level) if distributed: teacher_model = DataParallel(teacher_model, device_ids=device_ids) student_model = DistributedDataParallel(student_model, device_ids=device_ids) start_epoch = args.start_epoch if not args.test_only: distill(teacher_model, student_model, train_data_loader, val_data_loader, device, distributed, start_epoch, config, args) student_model_without_ddp =\ student_model.module if isinstance(student_model, DistributedDataParallel) else student_model load_ckpt(student_model_config['ckpt'], model=student_model_without_ddp, strict=True) if not args.student_only: evaluate(teacher_model, test_data_loader, device, title='[Teacher: {}]'.format(teacher_model_type)) evaluate(student_model, test_data_loader, device, title='[Student: {}]'.format(student_model_config['type']))