def validate_one_epoch(self, epoch, verbose=True):
"""Evaluate on validation set"""
self.load_datasets()
with collectors_context(
self.activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(self.val_loader, self.model,
self.criterion, [self.pylogger],
self.args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[self.tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
if verbose:
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=[self.tflogger])
return top1, top5, vloss
def train_one_epoch(self, epoch, verbose=True):
"""Train for one epoch"""
self.load_datasets()
with collectors_context(
self.activations_collectors["train"]) as collectors:
top1, top5, loss = train(self.train_loader,
self.model,
self.criterion,
self.optimizer,
epoch,
self.compression_scheduler,
loggers=[self.tflogger, self.pylogger],
args=self.args)
if verbose:
distiller.log_weights_sparsity(self.model, epoch,
[self.tflogger, self.pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[self.tflogger],
collector=collectors["sparsity"])
if self.args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
self.model, self.compression_scheduler))
return top1, top5, loss
def test(test_loader, model, criterion, loggers, activations_collectors, args):
"""Model Test"""
msglogger.info('--- test ---------------------')
with collectors_context(activations_collectors["test"]) as collectors:
top1, top5, lossses = _validate(test_loader, model, criterion, loggers, args)
distiller.log_activation_statsitics(-1, "test", loggers, collector=collectors['sparsity'])
return top1, top5, lossses
def test(test_loader, model, criterion, loggers, activations_collectors, msglogger, args):
"""Model Test"""
if activations_collectors is None:
activations_collectors = create_activation_stats_collectors(model, None)
with collectors_context(activations_collectors["test"]) as collectors:
lossses, top1, top5, = val(model, criterion, test_loader, msglogger=msglogger)
distiller.log_activation_statsitics(-1, "test", loggers, collector=collectors['sparsity'])
return top1, top5, lossses
def test(test_loader, model, criterion, loggers, activations_collectors, args):
"""Model Test"""
msglogger.info('--- test ---------------------')
if activations_collectors is None:
activations_collectors = create_activation_stats_collectors(
model, None)
with collectors_context(activations_collectors["test"]) as collectors:
args.validset_print_period = args.testset_print_period
top1, top5, lossses = _validate(test_loader, model, criterion, loggers,
args)
distiller.log_activation_statsitics(-1,
"test",
loggers,
collector=collectors['sparsity'])
save_collectors_data(collectors, msglogger.logdir)
return top1, top5, lossses
def test(test_loader, model, criterion, loggers, activations_collectors, args):
"""Model Test"""
msglogger.info('--- test ---------------------')
if activations_collectors is None:
activations_collectors = create_activation_stats_collectors(
model, None)
with collectors_context(activations_collectors["test"]) as collectors:
df_column_entry_dict = {}
event = "validation"
if torch.cuda.is_available():
print_memory_metrics("before " + event, df_column_entry_dict)
start_mem_measurement()
start = time.time()
top1, top5, lossses = _validate(test_loader, model, criterion, loggers,
args)
time_elapse = time.time() - start
df_column_entry_dict['Time measurement at ' + event +
' [s]'] = time_elapse
stop_mem_measurement(event, df_column_entry_dict)
if torch.cuda.is_available():
print_memory_metrics("after " + event, df_column_entry_dict)
import pandas as pd
frame = pd.DataFrame([list(df_column_entry_dict.values())],
columns=list(df_column_entry_dict.keys()))
frame.to_excel(DISTILLER_PATH + msglogger.logdir + ".xlsx",
index=False)
distiller.log_activation_statsitics(-1,
"test",
loggers,
collector=collectors['sparsity'])
save_collectors_data(collectors, msglogger.logdir)
return top1, top5, lossses
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.evaluate:
args.deterministic = True
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/
distiller.set_deterministic() # Use a well-known seed, for repeatability of experiments
else:
# 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
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 'cinic' in args.arch:
args.dataset = 'cinic10'
else:
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset in ['cifar10', 'cinic10'] 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)
model = create_model(False, args.dataset, args.arch, device_ids=args.gpus) # Get arch state_dict
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)
# Load Pre-trained Model
chkpt_file=args.resume
print("=> loading checkpoint %s" % chkpt_file)
checkpoint = torch.load(chkpt_file)
model.load_state_dict(checkpoint['net'])
# 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:
return sensitivity_analysis(model, criterion, test_loader, pylogger, args)
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)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.cuda()
else:
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)
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
#sparsity = distiller.model_sparsity(model)
is_best = top1 > best_epochs[0].top1
if is_best:
best_epochs[0].epoch = epoch
best_epochs[0].top1 = top1
#best_epoch.sparsity = sparsity
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[0].top1, is_best, args.name, msglogger.logdir)
# Finally run results on the test set
test(test_loader, model, criterion, [pylogger], activations_collectors, args=args)
def main(args):
utils.init_distributed_mode(args)
print(args)
device = torch.device(args.device)
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
if utils.is_main_process():
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir,
args.verbose)
# 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)
msglogger.debug("Distiller: %s", distiller.__version__)
else:
msglogger = logging.getLogger()
msglogger.disabled = True
# Data loading code
print("Loading data")
dataset, num_classes = get_dataset(args.dataset, "train", get_transform(train=True), args.data_path)
dataset_test, _ = get_dataset(args.dataset, "val", get_transform(train=False), args.data_path)
print("Creating data loaders")
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)
else:
train_sampler = torch.utils.data.RandomSampler(dataset)
test_sampler = torch.utils.data.SequentialSampler(dataset_test)
if args.aspect_ratio_group_factor >= 0:
group_ids = create_aspect_ratio_groups(dataset, k=args.aspect_ratio_group_factor)
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)
else:
train_batch_sampler = torch.utils.data.BatchSampler(
train_sampler, args.batch_size, drop_last=True)
data_loader = torch.utils.data.DataLoader(
dataset, batch_sampler=train_batch_sampler, num_workers=args.workers,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1,
sampler=test_sampler, num_workers=args.workers,
collate_fn=utils.collate_fn)
print("Creating model")
model = detection.__dict__[args.model](num_classes=num_classes,
pretrained=args.pretrained)
patch_fastrcnn(model)
model.to(device)
if args.summary:
if utils.is_main_process():
for summary in args.summary:
distiller.model_summary(model, summary, args.dataset)
return
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
model_without_ddp = model.module
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
compression_scheduler = None
if utils.is_main_process():
# 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)
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, 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.qe_calibration:
def test_fn(model):
return evaluate(model, data_loader_test, device=device)
collect_quant_stats(model_without_ddp, test_fn, save_dir=args.output_dir,
modules_to_collect=['backbone', 'rpn', 'roi_heads'])
# We skip `.transform` because it is a pre-processing unit.
return
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
if compression_scheduler and 'compression_scheduler' in checkpoint:
compression_scheduler.load_state_dict(checkpoint['compression_scheduler'])
if args.test_only:
evaluate(model, data_loader_test, device=device)
return
activations_collectors = create_activation_stats_collectors(model, *args.activation_stats)
print("Start training")
start_time = time.time()
# if not isinstance(model, nn.DataParallel) and torch.cuda.is_available() \
# and torch.cuda.device_count() > 1:
# msglogger.info("Using %d GPUs on DataParallel." % torch.cuda.device_count())
# model = nn.DataParallel(model)
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
dist.barrier()
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
with collectors_context(activations_collectors["train"]) as collectors:
train_one_epoch(model, optimizer, data_loader, device, epoch, args.print_freq, compression_scheduler)
if utils.is_main_process():
distiller.log_weights_sparsity(model, epoch, loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(epoch, "train", loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity and utils.is_main_process():
msglogger.info(distiller.masks_sparsity_tbl_summary(model, compression_scheduler))
lr_scheduler.step()
if args.output_dir:
save_dict = {
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'args': args}
if compression_scheduler:
save_dict['compression_scheduler'] = compression_scheduler.state_dict()
utils.save_on_master(save_dict,
os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
# evaluate after every epoch
evaluate(model, data_loader_test, device=device)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main():
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,
args.verbose)
# 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, config = 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)
if "ssd" in args.arch:
neg_pos_ratio = 3
criterion = MultiboxLoss(config.priors,
iou_threshold=0.5,
neg_pos_ratio=neg_pos_ratio,
center_variance=0.1,
size_variance=0.2,
device=args.device,
reduction="sum",
class_reduction=True,
verbose=0)
else:
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is None:
if "ssd" in args.arch:
base_net_lr = args.lr
extra_layers_lr = args.lr
params = [{
'params': model.base_net.parameters(),
'lr': base_net_lr
}, {
'params':
itertools.chain(model.source_layer_add_ons.parameters(),
model.extras.parameters()),
'lr':
extra_layers_lr
}, {
'params':
itertools.chain(model.regression_headers.parameters(),
model.classification_headers.parameters())
}]
else:
params = model.parameters()
optimizer = torch.optim.SGD(params,
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:
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, config=config)
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,
parallel=not args.load_serialized,
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)
raw_teacher_model_path = msglogger.logdir + "/raw_teacher.pth.tar"
if not os.path.exists(raw_teacher_model_path):
teacher.save(raw_teacher_model_path)
msglogger.info(Fore.CYAN + '\tRaw Teacher Model saved: {0}'.format(
raw_teacher_model_path) + Style.RESET_ALL)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model,
teacher,
args.kd_temp,
dlw,
loss_type=args.kd_loss_type,
focal_alpha=args.kd_focal_alpha,
use_adaptive=args.kd_focal_adaptive,
verbose=0)
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
}
try:
raw_fullpath_best = 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)
except Exception as ex:
# keep previous fullpath_best
pass
mlflow.log_artifacts(msglogger.logdir)
# Finally run results on the test set
eval_params = {
"model_type": args.arch,
"model_path": raw_fullpath_best,
"dataset_path": args.data,
"label_path": "models/voc-model-labels.txt"
}
mlflow.projects.run(uri=".",
entry_point="eval",
use_conda=False,
parameters=eval_params)
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)
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)
raise ValueError('Epochs parameter is too low. Nothing to do.')
print(start_epoch, ending_epoch)
for epoch in range(start_epoch, ending_epoch):
print(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,
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)
