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] = []