def val_segmentation(epoch, net_segmentation):
global best_iou
global val_seg_iou
progbar = tqdm(total=len(val_seg_loader), desc='Val')
net_segmentation.eval()
val_seg_loss.append(0)
hist = np.zeros((nClasses, nClasses))
for batch_idx, (inputs_, targets) in enumerate(val_seg_loader):
inputs_, targets = Variable(inputs_.to(device)), Variable(
targets.to(device))
outputs = net_segmentation(inputs_)
total_loss = 1 - soft_iou(outputs, targets, ignore=ignore_class)
val_seg_loss[-1] += total_loss.data
_, predicted = torch.max(outputs.data, 1)
correctLabel = targets.view(-1, targets.size()[1], targets.size()[2])
hist += fast_hist(
correctLabel.view(correctLabel.size(0), -1).cpu().numpy(),
predicted.view(predicted.size(0), -1).cpu().numpy(), nClasses)
miou, p_acc, fwacc = performMetrics(hist)
progbar.set_description('Val (loss=%.4f, mIoU=%.4f)' %
(val_seg_loss[-1] / (batch_idx + 1), miou))
progbar.update(1)
val_seg_loss[-1] = val_seg_loss[-1] / len(val_seg_loader)
val_miou, _, _ = performMetrics(hist)
val_seg_iou += [val_miou]
if best_iou < val_miou:
best_iou = val_miou
print('Saving..')
state = {'net_segmentation': net_segmentation}
torch.save(state,
model_root + experiment + 'segmentation' + '.ckpt.t7')
def train_segmentation(epoch, net_segmentation, seg_optimizer):
global train_seg_iou
progbar = tqdm(total=len(train_seg_loader), desc='Train')
net_segmentation.train()
train_seg_loss.append(0)
seg_optimizer.zero_grad()
hist = np.zeros((nClasses, nClasses))
for batch_idx, (inputs_, targets) in enumerate(train_seg_loader):
inputs_, targets = Variable(inputs_.to(device)), Variable(
targets.to(device))
outputs = net_segmentation(inputs_)
total_loss = (
1 - soft_iou(outputs, targets, ignore=ignore_class)) / ITER_SIZE
total_loss.backward()
if (batch_idx % ITER_SIZE == 0
and batch_idx != 0) or batch_idx == len(train_loader) - 1:
seg_optimizer.step()
seg_optimizer.zero_grad()
train_seg_loss[-1] += total_loss.data
_, predicted = torch.max(outputs.data, 1)
correctLabel = targets.view(-1, targets.size()[1], targets.size()[2])
hist += fast_hist(
correctLabel.view(correctLabel.size(0), -1).cpu().numpy(),
predicted.view(predicted.size(0), -1).cpu().numpy(), nClasses)
miou, p_acc, fwacc = performMetrics(hist)
progbar.set_description('Train (loss=%.4f, mIoU=%.4f)' %
(train_seg_loss[-1] / (batch_idx + 1), miou))
progbar.update(1)
train_seg_loss[-1] = train_seg_loss[-1] / len(train_seg_loader)
miou, p_acc, fwacc = performMetrics(hist)
train_seg_iou += [miou]
def evaluate_segmentation(net_segmentation):
net_segmentation.eval()
hist = np.zeros((nClasses, nClasses))
val_seg_loader = torch.utils.data.DataLoader(segmentation_data_loader(
img_root=val_img_root,
gt_root=val_gt_root,
image_list=val_image_list,
suffix=dataset,
out=out,
crop=False,
mirror=False),
batch_size=1,
num_workers=8,
shuffle=False)
progbar = tqdm(total=len(val_seg_loader), desc='Eval')
hist = np.zeros((nClasses, nClasses))
for batch_idx, (inputs_, targets) in enumerate(val_seg_loader):
inputs_, targets = Variable(inputs_.to(device)), Variable(
targets.to(device))
outputs = net_segmentation(inputs_)
_, predicted = torch.max(outputs.data, 1)
correctLabel = targets.view(-1, targets.size()[1], targets.size()[2])
hist += fast_hist(
correctLabel.view(correctLabel.size(0), -1).cpu().numpy(),
predicted.view(predicted.size(0), -1).cpu().numpy(), nClasses)
miou, p_acc, fwacc = performMetrics(hist)
progbar.set_description('Eval (mIoU=%.4f)' % (miou))
progbar.update(1)
miou, p_acc, fwacc = performMetrics(hist)
print('\n mIoU: ', miou)
print('\n Pixel accuracy: ', p_acc)
print('\n Frequency Weighted Pixel accuracy: ', fwacc)