def train(model: nn.Module,
loader: DataLoader,
class_loss: nn.Module,
optimizer: Optimizer,
scheduler: _LRScheduler,
epoch: int,
callback: VisdomLogger,
freq: int,
ex: Experiment = None) -> None:
model.train()
device = next(model.parameters()).device
to_device = lambda x: x.to(device, non_blocking=True)
loader_length = len(loader)
train_losses = AverageMeter(device=device, length=loader_length)
train_accs = AverageMeter(device=device, length=loader_length)
pbar = tqdm(loader, ncols=80, desc='Training [{:03d}]'.format(epoch))
for i, (batch, labels, indices) in enumerate(pbar):
batch, labels, indices = map(to_device, (batch, labels, indices))
logits, features = model(batch)
loss = class_loss(logits, labels).mean()
acc = (logits.detach().argmax(1) == labels).float().mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
train_losses.append(loss)
train_accs.append(acc)
if callback is not None and not (i + 1) % freq:
step = epoch + i / loader_length
callback.scalar('xent',
step,
train_losses.last_avg,
title='Train Losses')
callback.scalar('train_acc',
step,
train_accs.last_avg,
title='Train Acc')
if ex is not None:
for i, (loss, acc) in enumerate(
zip(train_losses.values_list, train_accs.values_list)):
step = epoch + i / loader_length
ex.log_scalar('train.loss', loss, step=step)
ex.log_scalar('train.acc', acc, step=step)
def __call__(self, ex: sacred.Experiment, mode, k, v):
if mode == 'train':
self.train_emas[k] = (
self.ema_beta * v +
(1.0 - self.ema_beta) * self.train_emas.get(k, v))
self.train_vals[k] = self.train_vals.get(k, []) + [v]
ex.log_scalar('training.{k}', self.train_emas[k])
elif mode == 'val':
ex.log_scalar('val.{k}', np.mean(np.array(v)))
ex.log_scalar('train.{k}', np.mean(np.array(self.train_vals[k])))
self.train_vals[k] = []
