def evaluation(model, criterion, loader, device, config, mode='val'):
y_true = []
y_pred = []
acc_meter = tnt.meter.ClassErrorMeter(accuracy=True)
loss_meter = tnt.meter.AverageValueMeter()
for (x, y) in loader:
y_true.extend(list(map(int, y)))
x = recursive_todevice(x, device)
y = y.to(device)
with torch.no_grad():
prediction = model(x)
loss = criterion(prediction, y)
acc_meter.add(prediction, y)
loss_meter.add(loss.item())
y_p = prediction.argmax(dim=1).cpu().numpy()
y_pred.extend(list(y_p))
metrics = {
'{}_accuracy'.format(mode): acc_meter.value()[0],
'{}_loss'.format(mode): loss_meter.value()[0],
'{}_IoU'.format(mode): mIou(y_true, y_pred, config['num_classes'])
}
if mode == 'val':
return metrics
elif mode == 'test':
return metrics, confusion_matrix(y_true,
y_pred,
labels=list(
range(config['num_classes'])))
def train_epoch(model, optimizer, criterion, data_loader, device, config):
acc_meter = tnt.meter.ClassErrorMeter(accuracy=True)
loss_meter = tnt.meter.AverageValueMeter()
y_true = []
y_pred = []
for i, (x, y) in enumerate(data_loader):
y_true.extend(list(map(int, y)))
x = recursive_todevice(x, device)
y = y.to(device)
optimizer.zero_grad()
out = model(x)
loss = criterion(out, y.long())
loss.backward()
optimizer.step()
pred = out.detach()
y_p = pred.argmax(dim=1).cpu().numpy()
y_pred.extend(list(y_p))
acc_meter.add(pred, y)
loss_meter.add(loss.item())
if (i + 1) % config['display_step'] == 0:
print('Step [{}/{}], Loss: {:.4f}, Acc : {:.2f}'.format(
i + 1, len(data_loader),
loss_meter.value()[0],
acc_meter.value()[0]))
epoch_metrics = {
'train_loss': loss_meter.value()[0],
'train_accuracy': acc_meter.value()[0],
'train_IoU': mIou(y_true, y_pred, n_classes=config['num_classes'])
}
return epoch_metrics