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