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']))
