def train(self, print_freq=10, fixbase_epoch=0, open_layers=None):
losses = MetricMeter()
batch_time = AverageMeter()
data_time = AverageMeter()
accuracy = AverageMeter()
self.set_model_mode('train')
self.two_stepped_transfer_learning(
self.epoch, fixbase_epoch, open_layers
)
self.num_batches = len(self.train_loader)
end = time.time()
for self.batch_idx, data in enumerate(self.train_loader):
data_time.update(time.time() - end)
loss_summary, avg_acc = self.forward_backward(data)
batch_time.update(time.time() - end)
losses.update(loss_summary)
accuracy.update(avg_acc)
if (self.batch_idx + 1) % print_freq == 0:
nb_this_epoch = self.num_batches - (self.batch_idx + 1)
nb_future_epochs = (self.max_epoch - (self.epoch + 1)) * self.num_batches
eta_seconds = batch_time.avg * (nb_this_epoch+nb_future_epochs)
eta_str = str(datetime.timedelta(seconds=int(eta_seconds)))
print(
'epoch: [{0}/{1}][{2}/{3}]\t'
'time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'cls acc {accuracy.val:.3f} ({accuracy.avg:.3f})\t'
'eta {eta}\t'
'{losses}\t'
'lr {lr:.6f}'.format(
self.epoch + 1,
self.max_epoch,
self.batch_idx + 1,
self.num_batches,
batch_time=batch_time,
data_time=data_time,
accuracy=accuracy,
eta=eta_str,
losses=losses,
lr=self.get_current_lr()
)
)
if self.writer is not None:
n_iter = self.epoch * self.num_batches + self.batch_idx
self.writer.add_scalar('Train/time', batch_time.avg, n_iter)
self.writer.add_scalar('Train/data', data_time.avg, n_iter)
self.writer.add_scalar('Aux/lr', self.get_current_lr(), n_iter)
self.writer.add_scalar('Accuracy/train', accuracy.avg, n_iter)
for name, meter in losses.meters.items():
self.writer.add_scalar('Loss/' + name, meter.avg, n_iter)
end = time.time()
self.update_lr()
def train(self,
print_freq=10,
fixbase_epoch=0,
open_layers=None,
lr_finder=False,
perf_monitor=None,
stop_callback=None):
losses = MetricMeter()
batch_time = AverageMeter()
data_time = AverageMeter()
accuracy = AverageMeter()
self.set_model_mode('train')
if not self._should_freeze_aux_models(self.epoch):
# NB: it should be done before `two_stepped_transfer_learning`
# to give possibility to freeze some layers in the unlikely event
# that `two_stepped_transfer_learning` is used together with nncf
self._unfreeze_aux_models()
self.two_stepped_transfer_learning(self.epoch, fixbase_epoch,
open_layers)
if self._should_freeze_aux_models(self.epoch):
self._freeze_aux_models()
self.num_batches = len(self.train_loader)
end = time.time()
for self.batch_idx, data in enumerate(self.train_loader):
if perf_monitor and not lr_finder:
perf_monitor.on_train_batch_begin(self.batch_idx)
data_time.update(time.time() - end)
if self.compression_ctrl:
self.compression_ctrl.scheduler.step(self.batch_idx)
loss_summary, avg_acc = self.forward_backward(data)
batch_time.update(time.time() - end)
losses.update(loss_summary)
accuracy.update(avg_acc)
if perf_monitor and not lr_finder:
perf_monitor.on_train_batch_end(self.batch_idx)
if not lr_finder and (((self.batch_idx + 1) % print_freq) == 0
or self.batch_idx == self.num_batches - 1):
nb_this_epoch = self.num_batches - (self.batch_idx + 1)
nb_future_epochs = (self.max_epoch -
(self.epoch + 1)) * self.num_batches
eta_seconds = batch_time.avg * (nb_this_epoch +
nb_future_epochs)
eta_str = str(datetime.timedelta(seconds=int(eta_seconds)))
print('epoch: [{0}/{1}][{2}/{3}]\t'
'time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'cls acc {accuracy.val:.3f} ({accuracy.avg:.3f})\t'
'eta {eta}\t'
'{losses}\t'
'lr {lr:.6f}'.format(self.epoch + 1,
self.max_epoch,
self.batch_idx + 1,
self.num_batches,
batch_time=batch_time,
data_time=data_time,
accuracy=accuracy,
eta=eta_str,
losses=losses,
lr=self.get_current_lr()))
if self.writer is not None and not lr_finder:
n_iter = self.epoch * self.num_batches + self.batch_idx
self.writer.add_scalar('Train/time', batch_time.avg, n_iter)
self.writer.add_scalar('Train/data', data_time.avg, n_iter)
self.writer.add_scalar('Aux/lr', self.get_current_lr(), n_iter)
self.writer.add_scalar('Accuracy/train', accuracy.avg, n_iter)
for name, meter in losses.meters.items():
self.writer.add_scalar('Loss/' + name, meter.avg, n_iter)
end = time.time()
self.current_lr = self.get_current_lr()
if stop_callback and stop_callback.check_stop():
break
if not lr_finder and self.use_ema_decay:
self.ema_model.update(self.models[self.main_model_name])
if self.per_batch_annealing:
self.update_lr()
return losses.meters['loss'].avg
