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_statistics(
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 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_statistics(
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 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:
top1, top5, lossses = _validate(test_loader, model, criterion, loggers, args)
distiller.log_activation_statistics(-1, "test", loggers, collector=collectors['sparsity'])
save_collectors_data(collectors, msglogger.logdir)
return top1, top5, lossses
