def validate(config, testloader, model, writer_dict):
model.eval()
ave_loss = AverageMeter()
confusion_matrix = np.zeros(
(config.DATASET.NUM_CLASSES, config.DATASET.NUM_CLASSES))
with torch.no_grad():
for _, batch in enumerate(testloader):
image, label, _, _ = batch
size = label.size()
label = label.long().cuda()
losses, pred = model(image, label)
pred = F.upsample(input=pred,
size=(size[-2], size[-1]),
mode='bilinear')
loss = losses.mean()
ave_loss.update(loss.item())
confusion_matrix += get_confusion_matrix(
label, pred, size, config.DATASET.NUM_CLASSES,
config.TRAIN.IGNORE_LABEL)
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
IoU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IoU = IoU_array.mean()
writer = writer_dict['writer']
global_steps = writer_dict['valid_global_steps']
writer.add_scalar('valid_loss', ave_loss.average(), global_steps)
writer.add_scalar('valid_mIoU', mean_IoU, global_steps)
writer_dict['valid_global_steps'] = global_steps + 1
return ave_loss.average(), mean_IoU, IoU_array
def train(config, epoch, num_epoch, epoch_iters, base_lr, num_iters,
trainloader, optimizer, model, writer_dict):
# Training
model.train()
batch_time = AverageMeter()
ave_loss = AverageMeter()
tic = time.time()
cur_iters = epoch * epoch_iters
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
for i_iter, batch in enumerate(trainloader, 0):
images, labels = batch
images = images.cuda()
labels = labels.long().cuda()
losses, _ = model(images, labels)
loss = losses.mean()
model.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - tic)
tic = time.time()
# update average loss
ave_loss.update(loss.item())
lr = adjust_learning_rate(optimizer, base_lr, num_iters,
i_iter + cur_iters)
if i_iter % config.PRINT_FREQ == 0:
msg = 'Epoch: [{}/{}] Iter:[{}/{}], Time: {:.2f}, ' \
'lr: {:.6f}, Loss: {:.6f}' .format(
epoch, num_epoch, i_iter, epoch_iters,
batch_time.average(), lr, ave_loss.average())
logging.info(msg)
writer.add_scalar('train_loss', ave_loss.average(), global_steps)
writer_dict['train_global_steps'] = global_steps + 1
