def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if args.epochs is None:
args.epochs = 90
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(
filter(None, [args.compress, args.qe_stats_file
]), # remove both None and empty strings
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
if args.evaluate:
args.deterministic = True
if args.deterministic:
distiller.set_deterministic(
args.seed) # For experiment reproducability
else:
if args.seed is not None:
distiller.set_seed(args.seed)
# Turn on CUDNN benchmark mode for best performance. This is usually "safe" for image
# classification models, as the input sizes don't change during the run
# See here: https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/3
cudnn.benchmark = True
start_epoch = 0
ending_epoch = args.epochs
perf_scores_history = []
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
raise ValueError(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
raise ValueError(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = distiller.apputils.classification_dataset_str_from_arch(
args.arch)
args.num_classes = distiller.apputils.classification_num_classes(
args.dataset)
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = create_model(args.pretrained,
args.dataset,
args.arch,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# TODO(barrh): args.deprecated_resume is deprecated since v0.3.1
if args.deprecated_resume:
msglogger.warning(
'The "--resume" flag is deprecated. Please use "--resume-from=YOUR_PATH" instead.'
)
if not args.reset_optimizer:
msglogger.warning(
'If you wish to also reset the optimizer, call with: --reset-optimizer'
)
args.reset_optimizer = True
args.resumed_checkpoint_path = args.deprecated_resume
# We can optionally resume from a checkpoint
optimizer = None
if args.resumed_checkpoint_path:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.resumed_checkpoint_path, model_device=args.device)
elif args.load_model_path:
model = apputils.load_lean_checkpoint(model,
args.load_model_path,
model_device=args.device)
if args.reset_optimizer:
start_epoch = 0
if optimizer is not None:
optimizer = None
msglogger.info(
'\nreset_optimizer flag set: Overriding resumed optimizer and resetting epoch count to 0'
)
# Define loss function (criterion)
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is None:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
for summary in args.summary:
distiller.model_summary(model, summary, args.dataset)
return
if args.export_onnx is not None:
return distiller.export_img_classifier_to_onnx(model,
os.path.join(
msglogger.logdir,
args.export_onnx),
args.dataset,
add_softmax=True,
verbose=False)
if args.qe_calibration:
return acts_quant_stats_collection(model, criterion, pylogger, args)
if args.activation_histograms:
return acts_histogram_collection(model, criterion, pylogger, args)
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = load_data(args)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, \
"You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resumed_checkpoint_path.replace(
".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
)
return
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_lean_checkpoint(teacher, args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy,
starting_epoch=args.kd_start_epoch,
ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'
.format(start_epoch, ending_epoch))
raise ValueError('Epochs parameter is too low. Nothing to do.')
for epoch in range(start_epoch, ending_epoch):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader,
model,
criterion,
optimizer,
epoch,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
if args.print_sparsity:
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Performance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch,
optimizer,
metrics={
'min': vloss,
'max': top1
})
# Update the list of top scores achieved so far, and save the checkpoint
update_training_scores_history(perf_scores_history, model, top1, top5,
epoch, args.num_best_scores)
is_best = epoch == perf_scores_history[0].epoch
checkpoint_extras = {
'current_top1': top1,
'best_top1': perf_scores_history[0].top1,
'best_epoch': perf_scores_history[0].epoch
}
apputils.save_checkpoint(epoch,
args.arch,
model,
optimizer=optimizer,
scheduler=compression_scheduler,
extras=checkpoint_extras,
is_best=is_best,
name=args.name,
dir=msglogger.logdir)
# ADDED: look at learning rate
msglogger.info(f'\n==> Learning rate: {get_lr(optimizer)}')
# Finally run results on the test set
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
def main():
global msglogger
check_pytorch_version()
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(sys.argv, gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
best_top1 = 0
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
msglogger.error('ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1')
exit(1)
# Use a well-known seed, for repeatability of experiments
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)
cudnn.deterministic = True
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error('ERROR: Argument --gpus must be a comma-separated list of integers only')
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error('ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
if args.dataset.lower() == 'none':
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
# Create the model
png_summary = args.summary is not None and args.summary.startswith('png')
is_parallel = not png_summary and args.summary != 'compute' # For PNG summary, parallel graphs are illegible
model = create_model(args.pretrained, args.dataset, args.arch, parallel=is_parallel, device_ids=args.gpus)
#If the teacher is not None, create teacher model and load it
if args.teacher_arch.lower() != 'none':
try:
teacher_model = create_model(True, args.dataset, args.teacher_arch, parallel=is_parallel,
device_ids=args.gpus)
print('The teacher model has been loaded from the torchvision pretrained model list. User input ignored')
except:
if args.teacher_weights.lower() == 'none':
raise ValueError('If you pass a teacher model you also have to pass the path of its weights')
teacher_model = create_model(False, args.dataset, args.teacher_arch, parallel=is_parallel,
device_ids=args.gpus)
try:
teacher_weights = torch.load(args.teacher_weights)
if 'state_dict' in teacher_weights:
teacher_weights = teacher_weights['state_dict']
teacher_model.load_state_dict(teacher_weights)
except:
raise ValueError('Unable to load teacher weights. Loading path {} resulted in error'.format(args.teacher_weights))
else:
teacher_model = None
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# We can optionally resume from a checkpoint
if args.resume:
model, compression_scheduler, start_epoch = apputils.load_checkpoint(
model, chkpt_file=args.resume)
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
# This sample application can be invoked to produce various summary reports.
if args.summary:
which_summary = args.summary
if which_summary.startswith('png'):
apputils.draw_img_classifier_to_file(model, 'model.png', args.dataset, which_summary == 'png_w_params')
else:
distiller.model_summary(model, which_summary, args.dataset)
exit()
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.workers, args.validation_size, args.deterministic)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler), len(test_loader.sampler))
if teacher_model is not None:
t1, t5, _ = _validate(test_loader, teacher_model, criterion, None, len(test_loader) + 10)
print('Teacher model passed in input has Top-1 of {} and Top-5 of {}'.format(t1, t5))
activations_sparsity = None
if args.activation_stats:
# If your model has ReLU layers, then those layers have sparse activations.
# ActivationSparsityCollector will collect information about this sparsity.
# WARNING! Enabling activation sparsity collection will significantly slow down training!
activations_sparsity = ActivationSparsityCollector(model)
if args.sensitivity is not None:
# This sample application can be invoked to execute Sensitivity Analysis on your
# model. The ouptut is saved to CSV and PNG.
msglogger.info("Running sensitivity tests")
test_fnc = partial(test, test_loader=test_loader, criterion=criterion,
loggers=[pylogger], print_freq=args.print_freq)
which_params = [param_name for param_name, _ in model.named_parameters()]
sensitivity = distiller.perform_sensitivity_analysis(model,
net_params=which_params,
sparsities=np.arange(0.0, 0.50, 0.05) if args.sensitivity == 'filter' else np.arange(0.0, 0.95, 0.05),
test_func=test_fnc,
group=args.sensitivity)
distiller.sensitivities_to_png(sensitivity, 'sensitivity.png')
distiller.sensitivities_to_csv(sensitivity, 'sensitivity.csv')
exit()
if args.evaluate:
# This sample application can be invoked to evaluate the accuracy of your model on
# the test dataset.
# You can optionally quantize the model to 8-bit integer before evaluation.
# For example:
# python3 compress_classifier.py --arch resnet20_cifar ../data.cifar10 -p=50 --resume=checkpoint.pth.tar --evaluate
if args.quantize:
model.cpu()
quantizer = quantization.SymmetricLinearQuantizer(model, 8, 8)
quantizer.prepare_model()
model.cuda()
top1, _, _ = test(test_loader, model, criterion, [pylogger], args.print_freq)
if args.quantize:
checkpoint_name = 'quantized'
apputils.save_checkpoint(0, args.arch, model, optimizer=None, best_top1=top1,
name='_'.split(args.name, checkpoint_name) if args.name else checkpoint_name,
dir=msglogger.logdir)
exit()
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(model, optimizer, args.compress)
for epoch in range(start_epoch, start_epoch + args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
if epoch >= args.start_distillation_from_epoch:
train(train_loader, model, criterion, optimizer, epoch, compression_scheduler,
loggers=[tflogger, pylogger], print_freq=args.print_freq, log_params_hist=args.log_params_histograms,
teacher_model=teacher_model, temperature_distillation=args.temp_distillation,
weight_distillation_loss=args.weight_distillation_loss)
else:
train(train_loader, model, criterion, optimizer, epoch, compression_scheduler,
loggers=[tflogger, pylogger], print_freq=args.print_freq, log_params_hist=args.log_params_histograms,
teacher_model=None)
distiller.log_weights_sparsity(model, epoch, loggers=[tflogger, pylogger])
if args.activation_stats:
distiller.log_activation_sparsity(epoch, loggers=[tflogger, pylogger],
collector=activations_sparsity)
# evaluate on validation set
top1, top5, vloss = validate(val_loader, model, criterion, [pylogger], args.print_freq, epoch)
stats = ('Peformance/Validation/',
OrderedDict([('Loss', vloss),
('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats, None, epoch, steps_completed=0, total_steps=1,
log_freq=1, loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# remember best top1 and save checkpoint
is_best = top1 > best_top1
best_top1 = max(top1, best_top1)
apputils.save_checkpoint(epoch, args.arch, model, optimizer, compression_scheduler, best_top1, is_best,
args.name, msglogger.logdir)
# Finally run results on the test set
test(test_loader, model, criterion, [pylogger], args.print_freq)
def main():
global msglogger
script_dir = os.path.dirname(__file__)
args = parse_args()
# Distiller loggers
msglogger = apputils.config_pylogger('logging.conf',
args.name,
output_dir=args.output_dir)
tflogger = TensorBoardLogger(msglogger.logdir)
# tflogger.log_gradients = True
# pylogger = PythonLogger(msglogger)
if args.seed is not None:
msglogger.info("Using seed = {}".format(args.seed))
torch.manual_seed(args.seed)
np.random.seed(seed=args.seed)
args.qe_mode = str(args.qe_mode).split('.')[1]
args.qe_clip_acts = str(args.qe_clip_acts).split('.')[1]
apputils.log_execution_env_state(sys.argv)
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
exit(1)
if len(args.gpus) > 1:
msglogger.error('ERROR: Only single GPU supported for NCF')
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Save configuration to file
config = {k: v for k, v in args.__dict__.items()}
config['timestamp'] = "{:.0f}".format(datetime.utcnow().timestamp())
config['local_timestamp'] = str(datetime.now())
run_dir = msglogger.logdir
msglogger.info("Saving config and results to {}".format(run_dir))
if not os.path.exists(run_dir) and run_dir != '':
os.makedirs(run_dir)
utils.save_config(config, run_dir)
# Check that GPUs are actually available
use_cuda = not args.no_cuda and torch.cuda.is_available()
t1 = time.time()
# Load Data
training = not (args.eval or args.qe_calibration
or args.activation_histograms)
msglogger.info('Loading data')
if training:
train_dataset = CFTrainDataset(
os.path.join(args.data, TRAIN_RATINGS_FILENAME),
args.negative_samples)
train_dataloader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
nb_users, nb_items = train_dataset.nb_users, train_dataset.nb_items
else:
train_dataset = None
train_dataloader = None
nb_users, nb_items = (138493, 26744)
test_ratings = load_test_ratings(
os.path.join(args.data, TEST_RATINGS_FILENAME)) # noqa: E501
test_negs = load_test_negs(os.path.join(args.data, TEST_NEG_FILENAME))
msglogger.info(
'Load data done [%.1f s]. #user=%d, #item=%d, #train=%s, #test=%d' %
(time.time() - t1, nb_users, nb_items,
str(train_dataset.mat.nnz) if training else 'N/A', len(test_ratings)))
# Create model
model = NeuMF(nb_users,
nb_items,
mf_dim=args.factors,
mf_reg=0.,
mlp_layer_sizes=args.layers,
mlp_layer_regs=[0. for i in args.layers],
split_final=args.split_final)
if use_cuda:
model = model.cuda()
msglogger.info(model)
msglogger.info("{} parameters".format(utils.count_parameters(model)))
# Save model text description
with open(os.path.join(run_dir, 'model.txt'), 'w') as file:
file.write(str(model))
compression_scheduler = None
start_epoch = 0
optimizer = None
if args.load:
if training:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.load)
if args.reset_optimizer:
start_epoch = 0
optimizer = None
else:
model = apputils.load_lean_checkpoint(model, args.load)
# Add loss to graph
criterion = nn.BCEWithLogitsLoss()
if use_cuda:
criterion = criterion.cuda()
if training and optimizer is None:
optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.compress:
compression_scheduler = distiller.file_config(model, optimizer,
args.compress)
model.cuda()
# Create files for tracking training
valid_results_file = os.path.join(run_dir, 'valid_results.csv')
if args.qe_calibration or args.activation_histograms:
calib = {
'portion':
args.qe_calibration,
'desc_str':
'quantization calibration stats',
'collect_func':
partial(distiller.data_loggers.collect_quant_stats,
inplace_runtime_check=True,
disable_inplace_attrs=True)
}
hists = {
'portion':
args.activation_histograms,
'desc_str':
'activation histograms',
'collect_func':
partial(distiller.data_loggers.collect_histograms,
activation_stats=None,
nbins=2048,
save_hist_imgs=True)
}
d = calib if args.qe_calibration else hists
distiller.utils.assign_layer_fq_names(model)
num_users = int(np.floor(len(test_ratings) * d['portion']))
msglogger.info(
"Generating {} based on {:.1%} of the test-set ({} users)".format(
d['desc_str'], d['portion'], num_users))
test_fn = partial(val_epoch,
ratings=test_ratings,
negs=test_negs,
K=args.topk,
use_cuda=use_cuda,
processes=args.processes,
num_users=num_users)
d['collect_func'](model=model,
test_fn=test_fn,
save_dir=run_dir,
classes=None)
return 0
if args.eval:
if args.quantize_eval and args.qe_calibration is None:
model.cpu()
quantizer = quantization.PostTrainLinearQuantizer.from_args(
model, args)
dummy_input = (torch.tensor([1]), torch.tensor([1]),
torch.tensor([True], dtype=torch.bool))
quantizer.prepare_model(dummy_input)
model.cuda()
distiller.utils.assign_layer_fq_names(model)
if args.eval_fp16:
model = model.half()
# Calculate initial Hit Ratio and NDCG
begin = time.time()
hits, ndcgs = val_epoch(model,
test_ratings,
test_negs,
args.topk,
use_cuda=use_cuda,
processes=args.processes)
val_time = time.time() - begin
hit_rate = np.mean(hits)
msglogger.info(
'Initial HR@{K} = {hit_rate:.4f}, NDCG@{K} = {ndcg:.4f}, val_time = {val_time:.2f}'
.format(K=args.topk,
hit_rate=hit_rate,
ndcg=np.mean(ndcgs),
val_time=val_time))
hit_rate = 0
if args.quantize_eval:
checkpoint_name = 'quantized'
apputils.save_checkpoint(0,
'NCF',
model,
optimizer=None,
extras={'quantized_hr@10': hit_rate},
name='_'.join([args.name, 'quantized'])
if args.name else checkpoint_name,
dir=msglogger.logdir)
return 0
total_samples = len(train_dataloader.sampler)
steps_per_epoch = math.ceil(total_samples / args.batch_size)
best_hit_rate = 0
best_epoch = 0
for epoch in range(start_epoch, args.epochs):
msglogger.info('')
model.train()
losses = utils.AverageMeter()
begin = time.time()
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch, optimizer)
loader = tqdm.tqdm(train_dataloader)
for batch_index, (user, item, label) in enumerate(loader):
user = torch.autograd.Variable(user, requires_grad=False)
item = torch.autograd.Variable(item, requires_grad=False)
label = torch.autograd.Variable(label, requires_grad=False)
if use_cuda:
user = user.cuda(async=True)
item = item.cuda(async=True)
label = label.cuda(async=True)
if compression_scheduler:
compression_scheduler.on_minibatch_begin(
epoch, batch_index, steps_per_epoch, optimizer)
outputs = model(user, item, torch.tensor([False],
dtype=torch.bool))
loss = criterion(outputs, label)
if compression_scheduler:
compression_scheduler.before_backward_pass(
epoch,
batch_index,
steps_per_epoch,
loss,
optimizer,
return_loss_components=False)
losses.update(loss.data.item(), user.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if compression_scheduler:
compression_scheduler.on_minibatch_end(epoch, batch_index,
steps_per_epoch,
optimizer)
# Save stats to file
description = (
'Epoch {} Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
epoch, loss=losses))
loader.set_description(description)
steps_completed = batch_index + 1
if steps_completed % args.log_freq == 0:
stats_dict = OrderedDict()
stats_dict['Loss'] = losses.avg
stats = ('Performance/Training/', stats_dict)
params = model.named_parameters(
) if args.log_params_histograms else None
distiller.log_training_progress(stats, params, epoch,
steps_completed,
steps_per_epoch, args.log_freq,
[tflogger])
tflogger.log_model_buffers(model,
['tracked_min', 'tracked_max'],
'Quant/Train/Acts/TrackedMinMax',
epoch, steps_completed,
steps_per_epoch, args.log_freq)
train_time = time.time() - begin
begin = time.time()
hits, ndcgs = val_epoch(model,
test_ratings,
test_negs,
args.topk,
use_cuda=use_cuda,
output=valid_results_file,
epoch=epoch,
processes=args.processes)
val_time = time.time() - begin
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
hit_rate = np.mean(hits)
mean_ndcgs = np.mean(ndcgs)
stats_dict = OrderedDict()
stats_dict['HR@{0}'.format(args.topk)] = hit_rate
stats_dict['NDCG@{0}'.format(args.topk)] = mean_ndcgs
stats = ('Performance/Validation/', stats_dict)
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
msglogger.info(
'Epoch {epoch}: HR@{K} = {hit_rate:.4f}, NDCG@{K} = {ndcg:.4f}, AvgTrainLoss = {loss.avg:.4f}, '
'train_time = {train_time:.2f}, val_time = {val_time:.2f}'.format(
epoch=epoch,
K=args.topk,
hit_rate=hit_rate,
ndcg=mean_ndcgs,
loss=losses,
train_time=train_time,
val_time=val_time))
is_best = False
if hit_rate > best_hit_rate:
best_hit_rate = hit_rate
is_best = True
best_epoch = epoch
extras = {
'current_hr@10': hit_rate,
'best_hr@10': best_hit_rate,
'best_epoch': best_epoch
}
apputils.save_checkpoint(epoch,
'NCF',
model,
optimizer,
compression_scheduler,
extras,
is_best,
dir=run_dir)
if args.threshold is not None:
if np.mean(hits) >= args.threshold:
msglogger.info("Hit threshold of {}".format(args.threshold))
break
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if args.epochs is None:
args.epochs = 90
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(args.compress,
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
ending_epoch = args.epochs
perf_scores_history = []
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
raise ValueError(
'ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1'
)
# Use a well-known seed, for repeatability of experiments
distiller.set_deterministic()
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
raise ValueError(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
raise ValueError(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset == 'cifar10' else 1000
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = create_model(args.pretrained,
args.dataset,
args.arch,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# TODO(barrh): args.deprecated_resume is deprecated since v0.3.1
if args.deprecated_resume:
msglogger.warning(
'The "--resume" flag is deprecated. Please use "--resume-from=YOUR_PATH" instead.'
)
if not args.reset_optimizer:
msglogger.warning(
'If you wish to also reset the optimizer, call with: --reset-optimizer'
)
args.reset_optimizer = True
args.resumed_checkpoint_path = args.deprecated_resume
# We can optionally resume from a checkpoint
optimizer = None
if args.resumed_checkpoint_path:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.resumed_checkpoint_path, model_device=args.device)
elif args.load_model_path:
model = apputils.load_lean_checkpoint(model,
args.load_model_path,
model_device=args.device)
if args.reset_optimizer:
start_epoch = 0
if optimizer is not None:
optimizer = None
msglogger.info(
'\nreset_optimizer flag set: Overriding resumed optimizer and resetting epoch count to 0'
)
# Define loss function (criterion)
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is None:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
if args.qe_calibration:
msglogger.info('Quantization calibration stats collection enabled:')
msglogger.info(
'\tStats will be collected for {:.1%} of test dataset'.format(
args.qe_calibration))
msglogger.info(
'\tSetting constant seeds and converting model to serialized execution'
)
distiller.set_deterministic()
model = distiller.make_non_parallel_copy(model)
activations_collectors.update(
create_quantization_stats_collector(model))
args.evaluate = True
args.effective_test_size = args.qe_calibration
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.workers, args.validation_split, args.deterministic,
args.effective_train_size, args.effective_valid_size,
args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, \
"You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resumed_checkpoint_path.replace(
".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
)
return
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_lean_checkpoint(teacher, args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy,
starting_epoch=args.kd_start_epoch,
ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'
.format(start_epoch, ending_epoch))
raise ValueError('Epochs parameter is too low. Nothing to do.')
for epoch in range(start_epoch, ending_epoch):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(
epoch, metrics=(vloss if (epoch != start_epoch) else 10**6))
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader,
model,
criterion,
optimizer,
epoch,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Performance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint
update_training_scores_history(perf_scores_history, model, top1, top5,
epoch, args.num_best_scores)
is_best = epoch == perf_scores_history[0].epoch
checkpoint_extras = {
'current_top1': top1,
'best_top1': perf_scores_history[0].top1,
'best_epoch': perf_scores_history[0].epoch
}
apputils.save_checkpoint(epoch,
args.arch,
model,
optimizer=optimizer,
scheduler=compression_scheduler,
extras=checkpoint_extras,
is_best=is_best,
name=args.name,
dir=msglogger.logdir)
# Finally run results on the test set
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
def main():
global msglogger
check_pytorch_version()
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(sys.argv, gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
best_epochs = [distiller.MutableNamedTuple({'epoch': 0, 'top1': 0, 'sparsity': 0})
for i in range(args.num_best_scores)]
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
msglogger.error('ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1')
exit(1)
# Use a well-known seed, for repeatability of experiments
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)
cudnn.deterministic = True
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error('ERROR: Argument --gpus must be a comma-separated list of integers only')
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error('ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
if 'cifar' in args.arch:
args.dataset = 'cifar10'
elif 'mnist' in args.arch:
args.dataset = 'mnist'
elif 'svhn' in args.arch:
args.dataset = 'svhn'
else:
args.dataset = 'imagenet'
args.num_classes = 10 if args.dataset == 'cifar10' or args.dataset == 'svhn' else 1000
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = create_model(args.pretrained, args.dataset, args.arch,
parallel=not args.load_serialized, device_ids=args.gpus)
#DEBUG
print('++++++++++++++++++++++++++++++++++++++++++')
print(model)
print('++++++++++++++++++++++++++++++++++++++++++')
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s', args.earlyexit_thresholds)
# We can optionally resume from a checkpoint
if args.resume:
model, compression_scheduler, start_epoch = apputils.load_checkpoint(
model, chkpt_file=args.resume)
model.cuda()
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.ADC:
return automated_deep_compression(model, criterion, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.workers, args.validation_size, args.deterministic)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler), len(test_loader.sampler))
activations_collectors = create_activation_stats_collectors(model, collection_phase=args.activation_stats)
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0], args.sensitivity_range[1], args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger, args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger, activations_collectors, args)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(model, optimizer, args.compress, compression_scheduler)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.cuda()
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained, args.dataset, args.kd_teacher, device_ids=args.gpus)
if args.kd_resume:
teacher, _, _ = apputils.load_checkpoint(teacher, chkpt_file=args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy, starting_epoch=args.kd_start_epoch, ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
for epoch in range(start_epoch, start_epoch + args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader, model, criterion, optimizer, epoch, compression_scheduler,
loggers=[tflogger, pylogger], args=args)
distiller.log_weights_sparsity(model, epoch, loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(epoch, "train", loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(distiller.masks_sparsity_tbl_summary(model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion, [pylogger], args, epoch)
distiller.log_activation_statsitics(epoch, "valid", loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
f= open("val_accuracy.txt","w+")
f.write(str(top1))
#f.write(str(top5))
f.close()
stats = ('Peformance/Validation/',
OrderedDict([('Loss', vloss),
('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats, None, epoch, steps_completed=0, total_steps=1, log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint
is_best = top1 > best_epochs[-1].top1
if top1 > best_epochs[0].top1:
best_epochs[0].epoch = epoch
best_epochs[0].top1 = top1
# Keep best_epochs sorted such that best_epochs[0] is the lowest top1 in the best_epochs list
best_epochs = sorted(best_epochs, key=lambda score: score.top1)
for score in reversed(best_epochs):
if score.top1 > 0:
msglogger.info('==> Best Top1: %.3f on Epoch: %d', score.top1, score.epoch)
apputils.save_checkpoint(epoch, args.arch, model, optimizer, compression_scheduler,
best_epochs[-1].top1, is_best, args.name, msglogger.logdir)
def main(opt):
# basic settings
os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.gpu_ids)[1:-1]
if torch.cuda.is_available():
device = "cuda"
torch.backends.cudnn.benchmark = True
else:
device = "cpu"
##################### Get Dataloader ####################
dataloader_train, dataloader_val = custom_get_dataloaders(opt)
# dummy_input is sample input of dataloaders
if hasattr(dataloader_val, "dataset"):
dummy_input = dataloader_val.dataset.__getitem__(0)
dummy_input = dummy_input[0]
dummy_input = dummy_input.unsqueeze(0)
else:
# for imagenet dali loader
dummy_input = torch.rand(1, 3, 224, 224)
##################### Create Baseline Model ####################
net = ModelWrapper(opt)
net.load_checkpoint(opt.checkpoint)
flops_before, params_before = model_summary(net.get_compress_part(),
dummy_input)
##################### Load Pruning Strategy ###############
compression_scheduler = distiller.file_config(net.get_compress_part(),
net.optimizer,
opt.compress_schedule_path)
channel_config = get_channel_config(opt.search_result,
opt.strategy_id) # pruning strategy
compression_scheduler = random_compression_scheduler(
compression_scheduler, channel_config)
###### Adaptive-BN-based Candidate Evaluation of Pruning Strategy ###
thinning(net, compression_scheduler, input_tensor=dummy_input)
flops_after, params_after = model_summary(net.get_compress_part(),
dummy_input)
ratio = flops_after / flops_before
print("FLOPs ratio:", ratio)
net = net.to(device)
net.parallel(opt.gpu_ids)
net.get_compress_part().train()
with torch.no_grad():
for index, sample in enumerate(tqdm(dataloader_train, leave=False)):
_ = net.get_loss(sample)
if index > 100:
break
strategy_score = net.get_eval_scores(dataloader_val)["accuracy"]
print("Result file:{}, Strategy ID:{}, Evaluation score:{}".format(
opt.search_result, opt.strategy_id, strategy_score))
##################### Fine-tuning #########################
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(
net.optimizer, opt.epoch)
reporter = Reporter(opt)
best_acc = 0
net._net.train()
for epoch in range(1, opt.epoch + 1):
reporter.log_metric("lr", net.optimizer.param_groups[0]["lr"], epoch)
train_loss = train_epoch(
net,
dataloader_train,
net.optimizer,
)
reporter.log_metric("train_loss", train_loss, epoch)
lr_scheduler.step()
scores = net.get_eval_scores(dataloader_val)
print("==> Evaluation: Epoch={} Acc={}".format(epoch, str(scores)))
reporter.log_metric("eval_acc", scores["accuracy"], epoch)
if scores["accuracy"] > best_acc:
best_acc = scores["accuracy"]
reporter.log_metric("best_acc", best_acc, epoch)
save_checkpoints(
scores["accuracy"],
net._net,
reporter,
opt.exp_name,
epoch,
)
print("==> Training epoch %d" % epoch)
def main():
global msglogger
check_pytorch_version()
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(sys.argv, gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
best_top1 = 0
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
msglogger.error('ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1')
exit(1)
# Use a well-known seed, for repeatability of experiments
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)
cudnn.deterministic = True
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error('ERROR: Argument --gpus must be a comma-separated list of integers only')
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error('ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset == 'cifar10' else 1000
# Create the model
model = create_model(args.pretrained, args.dataset, args.arch, device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# We can optionally resume from a checkpoint
if args.resume:
model, compression_scheduler, start_epoch = apputils.load_checkpoint(
model, chkpt_file=args.resume)
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.ADC:
return automated_deep_compression(model, criterion, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.workers, args.validation_size, args.deterministic)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler), len(test_loader.sampler))
activations_sparsity = None
if args.activation_stats:
# If your model has ReLU layers, then those layers have sparse activations.
# ActivationSparsityCollector will collect information about this sparsity.
# WARNING! Enabling activation sparsity collection will significantly slow down training!
activations_sparsity = ActivationSparsityCollector(model)
if args.sensitivity is not None:
return sensitivity_analysis(model, criterion, test_loader, pylogger, args)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger, args)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(model, optimizer, args.compress)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.cuda()
for epoch in range(start_epoch, start_epoch + args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
train(train_loader, model, criterion, optimizer, epoch, compression_scheduler,
loggers=[tflogger, pylogger], args=args)
distiller.log_weights_sparsity(model, epoch, loggers=[tflogger, pylogger])
if args.activation_stats:
distiller.log_activation_sparsity(epoch, loggers=[tflogger, pylogger],
collector=activations_sparsity)
# evaluate on validation set
top1, top5, vloss = validate(val_loader, model, criterion, [pylogger], args, epoch)
stats = ('Peformance/Validation/',
OrderedDict([('Loss', vloss),
('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats, None, epoch, steps_completed=0, total_steps=1,
log_freq=1, loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# remember best top1 and save checkpoint
is_best = top1 > best_top1
if is_best:
best_epoch = epoch
best_top1 = top1
msglogger.info('==> Best validation Top1: %.3f Epoch: %d', best_top1, best_epoch)
apputils.save_checkpoint(epoch, args.arch, model, optimizer, compression_scheduler, best_top1, is_best,
args.name, msglogger.logdir)
# Finally run results on the test set
test(test_loader, model, criterion, [pylogger], args=args)
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(args.compress,
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
best_epochs = list()
if args.deterministic:
if args.loaders is None:
args.loaders = 1
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.loaders > 1:
msglogger.error(
'ERROR: Setting --deterministic requires setting --loaders to 0 or 1'
)
exit(1)
# Use a well-known seed, for repeatability of experiments
distiller.set_deterministic()
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.use_cpu or (args.gpus is None
and not torch.cuda.is_available()) or (args.gpus
== ''):
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
if args.loaders is None:
active_gpus = args.gpus if args.gpus is not None else torch.cuda.device_count(
)
args.loaders = max(parser.DEFAULT_LOADERS_COUNT,
parser.DEFAULT_LOADERS_COUNT * active_gpus)
msglogger.debug('Number of data loaders set to: {}'.format(args.loaders))
# Infer the dataset from the model name
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset == 'cifar10' else 1000
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = create_model(args.pretrained,
args.dataset,
args.arch,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# We can optionally resume from a checkpoint
optimizer = None
resumed_training_steps = None
if args.resume or args.load_state_dict:
if args.resume and not args.reset_optimizer:
# initiate SGD with dummy lr
optimizer = torch.optim.SGD(model.parameters(), lr=0.36787944117)
model, compression_scheduler, optimizer, start_epoch, resumed_training_steps = apputils.load_checkpoint(
model, args.resume or args.load_state_dict, optimizer=optimizer)
model.to(args.device)
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is not None:
# optimizer was resumed from checkpoint
# check if user has tried to set optimizer arguments
# if so, ignore arguments with a warning.
optimizer_group_args = [
'lr', 'learning-rate', 'momentum', 'weight-decay', 'wd'
]
user_optim_args = [
x for x in optimizer_group_args for arg in sys.argv
if arg.startswith('--' + x)
]
if user_optim_args:
msglogger.warning(
'{} optimizer arguments are ignored.'.format(user_optim_args))
msglogger.info(
'setting optimizer arguments when optimizer is resumed '
'from checkpoint is forbidden.')
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
if args.qe_calibration:
msglogger.info('Quantization calibration stats collection enabled:')
msglogger.info(
'\tStats will be collected for {:.1%} of test dataset'.format(
args.qe_calibration))
msglogger.info(
'\tSetting constant seeds and converting model to serialized execution'
)
distiller.set_deterministic()
model = distiller.make_non_parallel_copy(model)
activations_collectors.update(
create_quantization_stats_collector(model))
args.evaluate = True
args.effective_test_size = args.qe_calibration
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.loaders, args.validation_split, args.deterministic,
args.effective_train_size, args.effective_valid_size,
args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
args.trainset_print_period = parser.getPrintPeriod(
args, len(train_loader.sampler), args.batch_size)
args.validset_print_period = parser.getPrintPeriod(args,
len(val_loader.sampler),
args.batch_size)
args.testset_print_period = parser.getPrintPeriod(args,
len(test_loader.sampler),
args.batch_size)
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if
(args.resume and not args.reset_optimizer) else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resume is not None, "You must use --resume to provide a checkpoint file to thinnify"
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resume.replace(".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
)
return
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_checkpoint(teacher,
chkpt_file=args.kd_resume)[0]
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(
args.kd_policy, range(args.kd_start_epoch, args.epochs, 1))
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if getattr(compression_scheduler, 'global_policy_end_epoch',
None) is not None:
if compression_scheduler.global_policy_end_epoch >= (start_epoch +
args.epochs):
msglogger.warning(
'scheduler requires at least {} epochs, but only {} are sanctioned'
.format(compression_scheduler.global_policy_end_epoch,
args.epochs))
accumulated_training_steps = resumed_training_steps if resumed_training_steps is not None else 0
for epoch in range(start_epoch, start_epoch + args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
try:
train(train_loader,
model,
criterion,
optimizer,
epoch,
accumulated_training_steps,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
except RuntimeError as e:
if ('cuda out of memory' in str(e).lower()):
msglogger.error(
'CUDA memory failure has been detected.\n'
'Sometimes it helps to decrease batch size.\n'
'e.g. Add the following flag to your call: --batch-size={}'
.format(args.batch_size // 10))
raise
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
accumulated_training_steps += math.ceil(
len(train_loader.sampler) / train_loader.batch_size)
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Performance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
tflogger.log_training_progress(stats, epoch, None)
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
if getattr(compression_scheduler, 'global_policy_end_epoch',
None) is None or (
compression_scheduler.global_policy_end_epoch <= epoch):
# Update the list of top scores achieved since all policies have concluded
if top1 > 0:
best_epochs.append(
distiller.MutableNamedTuple({
'top1': top1,
'top5': top5,
'epoch': epoch
}))
# Keep best_epochs sorted from best to worst
# Sort by top1 first, secondary sort by top5, and so forth
best_epochs.sort(key=operator.attrgetter('top1', 'top5', 'epoch'),
reverse=True)
for score in best_epochs[:args.num_best_scores]:
msglogger.info('==> Best Top1: %.3f Top5: %.3f on epoch: %d',
score.top1, score.top5, score.epoch)
is_best = best_epochs and (epoch == best_epochs[0].epoch)
apputils.save_checkpoint(epoch, args.arch, model, optimizer,
compression_scheduler,
best_epochs[0].top1 if best_epochs else None,
is_best, args.name, msglogger.logdir,
accumulated_training_steps)
# Finally run results on the test set
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
target_transform=gt_transform)
save_dir = os.path.join('pruning', model.name, pruner)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
stats_dir = os.path.join(save_dir, 'csv')
save_model_stats(model, stats_dir, 'initial')
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True)
writer = tensorboard.SummaryWriter(os.path.join(save_dir, 'logs'))
print('Dataset loaded successfully')
## Load disitller scheduler
compression_scheduler = distiller.file_config(model, optimizer, args.scheduler,
None, None)
for epoch in range(1, args.epoch + 1):
compression_scheduler.on_epoch_begin(
epoch) ## Call distiller on_epoch_begin
print('Started epoch {:03d}/{:03d}'.format(epoch, args.epoch))
train(train_loader, model, optimizer, epoch, writer, compression_scheduler)
compression_scheduler.on_epoch_end(epoch) ## Call distiller on_epoch_begin
torch.save(model.state_dict(),
'{}/{}.pruned.pth'.format(save_dir, model.name))
save_model_stats(model, stats_dir, epoch)
save_model_stats(model, stats_dir, 'final')
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if args.epochs is None:
args.epochs = 200
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(
filter(None, [args.compress, args.qe_stats_file
]), # remove both None and empty strings
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
if args.evaluate:
args.deterministic = True
if args.deterministic:
distiller.set_deterministic(
args.seed) # For experiment reproducability
else:
if args.seed is not None:
distiller.set_seed(args.seed)
# Turn on CUDNN benchmark mode for best performance. This is usually "safe" for image
# classification models, as the input sizes don't change during the run
# See here: https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/3
cudnn.benchmark = True
start_epoch = 0
ending_epoch = args.epochs
perf_scores_history = []
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
raise ValueError(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
raise ValueError(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
# TODO
args.dataset = 'coco'
# args.num_classes = 21 # wc -l ~/data/VOC2012/voc-model-labels.txt
if args.load_vgg19 and args.arch != 'vgg19':
raise ValueError(
'``load_vgg19`` should be set only when vgg19 is used')
model = create_pose_estimation_model(args.pretrained,
args.dataset,
args.arch,
load_vgg19=args.load_vgg19,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# <editor-fold desc=">>> Load Model">
# We can optionally resume from a checkpoint
optimizer = None
if args.resumed_checkpoint_path:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.resumed_checkpoint_path, model_device=args.device)
elif args.load_model_path:
model = apputils.load_lean_checkpoint(model,
args.load_model_path,
model_device=args.device)
if args.reset_optimizer:
start_epoch = 0
if optimizer is not None:
optimizer = None
msglogger.info(
'\nreset_optimizer flag set: Overriding resumed optimizer and resetting epoch count to 0'
)
# </editor-fold>
# Define loss function (criterion)
# get_loss(saved_for_loss, heat_temp, heat_weight,vec_temp, vec_weight)
criterion = {
'shufflenetv2': shufflenetv2_get_loss,
'vgg19': vgg19_get_loss,
'hourglass': hourglass_get_loss,
}[args.arch]
if optimizer is None:
trainable_vars = [
param for param in model.parameters() if param.requires_grad
]
optimizer = torch.optim.SGD(trainable_vars,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
# TODO: load lr_scheduler
lr_scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=0.8,
patience=5,
verbose=True,
threshold=0.0001,
threshold_mode='rel',
cooldown=3,
min_lr=0,
eps=1e-08)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
for summary in args.summary:
distiller.model_summary(model, summary, args.dataset)
return
if args.export_onnx is not None:
return distiller.export_img_classifier_to_onnx(model,
os.path.join(
msglogger.logdir,
args.export_onnx),
args.dataset,
add_softmax=True,
verbose=False)
if args.qe_calibration:
return acts_quant_stats_collection(model, criterion, pylogger, args)
if args.activation_histograms:
return acts_histogram_collection(model, criterion, pylogger, args)
print('Building activations_collectors...')
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
print('Loading data...')
train_loader, val_loader, test_loader, _ = load_data(args)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
# zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, \
"You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resumed_checkpoint_path.replace(
".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
)
return
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'
.format(start_epoch, ending_epoch))
raise ValueError('Epochs parameter is too low. Nothing to do.')
for epoch in range(start_epoch, ending_epoch):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(
epoch,
metrics=(total_loss if (epoch != start_epoch) else 10**6))
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader,
model,
criterion,
optimizer,
epoch,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
loss = validate(val_loader, model, criterion, [pylogger], args,
epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
lr_scheduler.step(loss)
stats = ('Performance/Validation/', OrderedDict([('Loss', loss)]))
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint
update_training_scores_history(perf_scores_history, model, loss, epoch,
args.num_best_scores)
is_best = epoch == perf_scores_history[0].epoch
checkpoint_extras = {
'current_loss': loss,
'best_loss': perf_scores_history[0].loss,
'best_epoch': perf_scores_history[0].epoch
}
apputils.save_checkpoint(epoch,
args.arch,
model,
optimizer=optimizer,
scheduler=compression_scheduler,
extras=checkpoint_extras,
is_best=is_best,
name=args.name,
dir=msglogger.logdir)
# Finally run results on the test set
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
def _init_learner(args):
# Create the model
model = create_model(args.pretrained,
args.dataset,
args.arch,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# TODO(barrh): args.deprecated_resume is deprecated since v0.3.1
if args.deprecated_resume:
msglogger.warning(
'The "--resume" flag is deprecated. Please use "--resume-from=YOUR_PATH" instead.'
)
if not args.reset_optimizer:
msglogger.warning(
'If you wish to also reset the optimizer, call with: --reset-optimizer'
)
args.reset_optimizer = True
args.resumed_checkpoint_path = args.deprecated_resume
optimizer = None
start_epoch = 0
if args.resumed_checkpoint_path:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.resumed_checkpoint_path, model_device=args.device)
elif args.load_model_path:
model = apputils.load_lean_checkpoint(model,
args.load_model_path,
model_device=args.device)
if args.reset_optimizer:
start_epoch = 0
if optimizer is not None:
optimizer = None
msglogger.info(
'\nreset_optimizer flag set: Overriding resumed optimizer and resetting epoch count to 0'
)
if optimizer is None:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.debug('Optimizer Type: %s', type(optimizer))
msglogger.debug('Optimizer Args: %s', optimizer.defaults)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
ending_epoch = args.epochs
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'
.format(start_epoch, ending_epoch))
raise ValueError('Epochs parameter is too low. Nothing to do.')
return model, compression_scheduler, optimizer, start_epoch, ending_epoch
args.effective_train_size, args.effective_valid_size, args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler), len(test_loader.sampler))
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0], args.sensitivity_range[1], args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger, args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger, activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(model, optimizer, args.compress, compression_scheduler,
(start_epoch-1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, "You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model, compression_scheduler.zeros_mask_dict, args.arch, args.dataset, optimizer=None)
apputils.save_checkpoint(0, args.arch, model, optimizer=None, scheduler=compression_scheduler,
name="{}_thinned".format(args.resumed_checkpoint_path.replace(".pth.tar", "")), dir=msglogger.logdir)
print("Note: your model may have collapsed to random inference, so you may want to fine-tune")
return
args.kd_policy = None
def main(opt, channel_config, dataloader_train, dataloader_val, path):
# basic settings
torch.backends.cudnn.enabled = False
os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.gpu_ids)[1:-1]
if torch.cuda.is_available():
device = "cuda"
torch.backends.cudnn.benchmark = False
else:
device = "cpu"
##################### Get Dataloader ####################
# dummy_input is sample input of dataloaders
if hasattr(dataloader_val, "dataset"):
dummy_input = dataloader_val.dataset.__getitem__(0)
dummy_input = dummy_input[0]
dummy_input = dummy_input.unsqueeze(0)
else:
# for imagenet dali loader
dummy_input = torch.rand(1, 3, 224, 224)
##################### Create Baseline Model ####################
net = ModelWrapper(opt)
load(net, path)
#net.load_checkpoint(opt.checkpoint)
##################### Load Pruning Strategy ###############
compression_scheduler = distiller.file_config(net.get_compress_part(),
net.optimizer,
opt.compress_schedule_path)
compression_scheduler = setCompressionScheduler(compression_scheduler,
channel_config)
###### Adaptive-BN-based Candidate Evaluation of Pruning Strategy ###
thinning(net, compression_scheduler, input_tensor=dummy_input)
flops_after, params_after = model_summary(net.get_compress_part(),
dummy_input)
net = net.to(device)
net.parallel(opt.gpu_ids)
net.get_compress_part().train()
t = tqdm(dataloader_train, leave=False)
with torch.no_grad():
for index, sample in enumerate(t):
_ = net.get_loss(sample)
if index > 100:
break
strategy_score = net.get_eval_scores(dataloader_val)["accuracy"]
old = strategy_score
print("Evaluation score:{}".format(strategy_score))
##################### Fine-tuning #########################
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(net.optimizer,
100,
eta_min=5e-5)
#lr_scheduler=optim.lr_scheduler.StepLR(net.optimizer,5,0.9)
reporter = Reporter(opt)
best_acc = strategy_score
best_kappa = 0
net._net.train()
for epoch in range(1, opt.epoch + 1):
net.confusion_matrix.reset()
reporter.log_metric("lr", net.optimizer.param_groups[0]["lr"], epoch)
train_loss = train_epoch(
net,
dataloader_train,
net.optimizer,
)
reporter.log_metric("train_loss", train_loss, epoch)
lr_scheduler.step()
scores = net.get_eval_scores(dataloader_val)
kappa = CaluKappa(net.confusion_matrix)
print("==> Evaluation: Epoch={} Acc={}".format(epoch, str(scores)))
reporter.log_metric("eval_acc", scores["accuracy"], epoch)
reporter.log_metric("kappa", kappa, epoch)
if scores["accuracy"] > best_acc:
best_acc = scores["accuracy"]
best_kappa = kappa
save_checkpoints(
scores["accuracy"],
net._net,
reporter,
opt.exp_name,
epoch,
)
reporter.log_metric("best_acc", best_acc, epoch)
save_checkpoints(
scores["accuracy"],
net._net,
reporter,
opt.exp_name,
epoch,
)
print("==> Training epoch %d" % epoch)
"""将模型转换为torch script保存"""
ckpt_name = "{}_best.pth".format(opt.exp_name)
load(net, os.path.join(reporter.ckpt_log_dir, ckpt_name))
net._net.eval()
traced_script_module = torch.jit.trace(net._net,
torch.rand(1, 3, 256, 256))
traced_script_module.save(os.path.join(reporter.log_dir, "model.pt"))
del net
return old, best_acc, best_kappa, flops_after, params_after
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
# 记录有关执行环境的各种详细信息。有时是有用的
# 参考过去的实验执行，这些信息可能有用。
apputils.log_execution_env_state(args.compress,
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
perf_scores_history = []
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
msglogger.error(
'ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1'
) # 错误：设置--确定性要求将--workers/-j设置为0或1
exit(1) # 正常退出程序
# Use a well-known seed, for repeatability of experiments 使用一种众所周知的种子，用于实验的重复性。
distiller.set_deterministic()
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = 'cousm'
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = ResNet152()
# model = torch.nn.DataParallel(model, device_ids=args.gpus) # 并行GPU
model.to(args.device)
compression_scheduler = None # 压缩调度
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
# 创建两个日志后端 TensorBoardLogger以Google的Tensor板可以读取的格式写入日志文件。python logger将写入python记录器。
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# We can optionally resume from a checkpoint
if args.resume: # 加载训练模型
# checkpoint = torch.load(args.resume)
# model.load_state_dict(checkpoint['state_dict'])
model, compression_scheduler, start_epoch = apputils.load_checkpoint(
model, chkpt_file=args.resume)
model.to(args.device)
# Define loss function (criterion) and optimizer # 定义损失函数和优化器SGD
criterion = nn.CrossEntropyLoss().to(args.device)
# optimizer = torch.optim.SGD(model.fc.parameters(), lr=args.lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
optimizer = torch.optim.Adam(model.model.fc.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC: # 自动化的深层压缩
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy: # 贪婪的
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports. # 可以调用此示例应用程序来生成各种摘要报告。
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
# 激活统计收集器
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
if args.qe_calibration:
msglogger.info('Quantization calibration stats collection enabled:')
msglogger.info(
'\tStats will be collected for {:.1%} of test dataset'.format(
args.qe_calibration))
msglogger.info(
'\tSetting constant seeds and converting model to serialized execution'
)
distiller.set_deterministic()
model = distiller.make_non_parallel_copy(model)
activations_collectors.update(
create_quantization_stats_collector(model)) # 量化统计收集器
args.evaluate = True
args.effective_test_size = args.qe_calibration
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
# 加载数据集：从传递的模型名称推断要加载的数据集
train_loader, val_loader, test_loader, _ = get_data_loaders(
datasets_fn, r'/home/tian/Desktop/image_yasuo', args.batch_size,
args.workers, args.validation_split, args.deterministic,
args.effective_train_size, args.effective_valid_size,
args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
# 可以调用此示例应用程序来对模型执行敏感性分析。输出保存到csv和png。
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
# #这个示例应用程序的主要用例是CNN压缩
# #需要yaml中的压缩计划配置文件。
compression_scheduler = distiller.file_config(model, optimizer,
args.compress,
compression_scheduler)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
# 如果添加了参数（如PactQualifier），则模型会重新传输到GPU。
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model) # 压缩计划程序
if args.thinnify:
# zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resume is not None, "You must use --resume to provide a checkpoint file to thinnify" # 必须使用--resume提供检查点文件以细化
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resume.replace(".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
) # 注意：您的模型可能已折叠为随机推理，因此您可能需要对其进行微调。
return
args.kd_policy = None # 蒸馏
if args.kd_teacher:
teacher = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
device_ids=args.gpus)
if args.kd_resume:
teacher, _, _ = apputils.load_checkpoint(teacher,
chkpt_file=args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy,
starting_epoch=args.kd_start_epoch,
ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
lr = args.lr
lr_decay = 0.5
for epoch in range(start_epoch, args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader,
model,
criterion,
optimizer,
epoch,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity: # 打印掩盖稀疏表 在end of each epoch
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Peformance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint # 更新到目前为止获得的最高分数列表，并保存检查点
sparsity = distiller.model_sparsity(model)
perf_scores_history.append(
distiller.MutableNamedTuple({
'sparsity': sparsity,
'top1': top1,
'top5': top5,
'epoch': epoch
}))
# Keep perf_scores_history sorted from best to worst
# Sort by sparsity as main sort key, then sort by top1, top5 and epoch
# 保持绩效分数历史记录从最好到最差的排序
# 按稀疏度排序为主排序键，然后按top1、top5、epoch排序
perf_scores_history.sort(key=operator.attrgetter(
'sparsity', 'top1', 'top5', 'epoch'),
reverse=True)
for score in perf_scores_history[:args.num_best_scores]:
msglogger.info(
'==> Best [Top1: %.3f Top5: %.3f Sparsity: %.2f on epoch: %d]',
score.top1, score.top5, score.sparsity, score.epoch)
is_best = epoch == perf_scores_history[0].epoch
apputils.save_checkpoint(epoch, args.arch, model, optimizer,
compression_scheduler,
perf_scores_history[0].top1, is_best,
args.name, msglogger.logdir)
if not is_best:
lr = lr * lr_decay
# 当loss大于上一次loss,降低学习率
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Finally run results on the test set # 最后在测试集上运行结果
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
