def train(train_loader): pbar = ProgressBar(n_batch=len(train_loader)) train_loss = AverageMeter() train_acc = AverageMeter() count = 0 model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output, loss = model(data, y=target, loss_fn=nn.CrossEntropyLoss()) pred = output.argmax( dim=1, keepdim=True) # get the index of the max log-probability correct = pred.eq(target.view_as(pred)).sum().item() loss.backward() optimizer.step() count += data.size(0) train_acc.update(correct, n=1) pbar.batch_step(batch_idx=batch_idx, info={ 'loss': loss.item(), 'acc': correct / data.size(0) }, bar_type='Training') train_loss.update(loss.item(), n=1) print(' ') return {'loss': train_loss.avg, 'acc': train_acc.sum / count}
def train(train_loader): pbar = ProgressBar(n_batch=len(train_loader)) train_loss = AverageMeter() model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = loss_fn(output, target) loss.backward() optimizer.step() pbar.batch_step(batch_idx=batch_idx, info={'loss': loss.item()}, bar_type='Training') train_loss.update(loss.item(), n=1) return {'loss': train_loss.avg}
def test(test_loader): pbar = ProgressBar(n_batch=len(test_loader)) valid_loss = AverageMeter() valid_acc = AverageMeter() model.eval() count = 0 with torch.no_grad(): for batch_idx, (data, target) in enumerate(test_loader): data, target = data.to(device), target.to(device) output, loss = model(data, y=target, loss_fn=nn.CrossEntropyLoss()) pred = output.argmax( dim=1, keepdim=True) # get the index of the max log-probability correct = pred.eq(target.view_as(pred)).sum().item() valid_loss.update(loss, n=data.size(0)) valid_acc.update(correct, n=1) count += data.size(0) pbar.batch_step(batch_idx=batch_idx, info={}, bar_type='Testing') return {'valid_loss': valid_loss.avg, 'valid_acc': valid_acc.sum / count}