def validate(epoch, val_loader, net, criterion, restore, best):
net.eval()
batch_inputs = []
batch_outputs = []
batch_labels = []
for vi, data in enumerate(val_loader, 0):
inputs, labels = data
inputs = Variable(inputs, volatile=True).cuda()
labels = Variable(labels, volatile=True).cuda()
outputs = net(inputs)
batch_inputs.append(inputs.cpu())
batch_outputs.append(outputs.cpu())
batch_labels.append(labels.cpu())
batch_inputs = torch.cat(batch_inputs)
batch_outputs = torch.cat(batch_outputs)
batch_labels = torch.cat(batch_labels)
val_loss = criterion(batch_outputs, batch_labels)
val_loss = val_loss.data[0]
batch_inputs = batch_inputs.data
batch_outputs = batch_outputs.data
batch_labels = batch_labels.data.numpy()
batch_prediction = batch_outputs.max(1)[1].squeeze_(1).numpy()
mean_iu = calculate_mean_iu(batch_prediction, batch_labels, num_classes)
if val_loss < best[0]:
best[0] = val_loss
best[1] = mean_iu
best[2] = epoch
torch.save(
net.state_dict(),
os.path.join(
ckpt_path, 'epoch_%d_validation_loss_%.4f_mean_iu_%.4f.pth' %
(epoch + 1, val_loss, mean_iu)))
to_save_dir = os.path.join(ckpt_path, str(epoch + 1))
rmrf_mkdir(to_save_dir)
for idx, tensor in enumerate(
zip(batch_inputs, batch_prediction, batch_labels)):
pil_input = restore(tensor[0])
pil_output = Image.fromarray(
colorize_mask(tensor[1], ignored_label=ignored_label), 'RGB')
pil_label = Image.fromarray(
colorize_mask(tensor[2], ignored_label=ignored_label), 'RGB')
pil_input.save(os.path.join(to_save_dir, '%d_img.png' % idx))
pil_output.save(os.path.join(to_save_dir, '%d_out.png' % idx))
pil_label.save(os.path.join(to_save_dir, '%d_label.png' % idx))
print '--------------------------------------------------------'
print '[validation loss %.4f]' % val_loss
print '[best validation loss %.4f], [best_mean_iu %.4f], [best epoch %d]' % (
best[0], best[1], best[2] + 1)
print '--------------------------------------------------------'
net.train()
def train(train_loader, net, criterion, optimizer, epoch):
for i, data in enumerate(train_loader, 0):
inputs, labels = data
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if (i + 1) % train_args['print_freq'] == 0:
outputs = outputs[:, :num_classes - 1, :, :]
prediction = outputs.data.max(1)[1].squeeze_(1).cpu().numpy()
mean_iu = calculate_mean_iu(prediction, labels.data.cpu().numpy(), num_classes)
print '[epoch %d], [iter %d], [training loss %.4f], [mean_iu %.4f]' % (
epoch + 1, i + 1, loss.data[0], mean_iu)
def validate(val_loader, net, criterion, optimizer, epoch, restore):
net.eval()
criterion.cpu()
input_batches = []
output_batches = []
label_batches = []
for vi, data in enumerate(val_loader, 0):
inputs, labels = data
inputs = Variable(inputs, volatile=True).cuda()
labels = Variable(labels, volatile=True).cuda()
outputs = net(inputs)
input_batches.append(inputs.cpu().data)
output_batches.append(outputs.cpu())
label_batches.append(labels.cpu())
input_batches = torch.cat(input_batches)
output_batches = torch.cat(output_batches)
label_batches = torch.cat(label_batches)
val_loss = criterion(output_batches, label_batches)
val_loss = val_loss.data[0]
output_batches = output_batches.cpu().data[:, :num_classes - 1, :, :]
label_batches = label_batches.cpu().data.numpy()
prediction_batches = output_batches.max(1)[1].squeeze_(1).numpy()
mean_iu = calculate_mean_iu(prediction_batches, label_batches, num_classes)
writer.add_scalar('loss', val_loss, epoch + 1)
writer.add_scalar('mean_iu', mean_iu, epoch + 1)
if val_loss < train_record['best_val_loss']:
train_record['best_val_loss'] = val_loss
train_record['corr_epoch'] = epoch + 1
train_record['corr_mean_iu'] = mean_iu
snapshot_name = 'epoch_%d_loss_%.4f_mean_iu_%.4f_lr_%.8f' % (
epoch + 1, val_loss, mean_iu, train_args['new_lr'])
torch.save(net.state_dict(),
os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
with open(exp_name + '.txt', 'a') as f:
f.write(snapshot_name + '\n')
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch + 1))
rmrf_mkdir(to_save_dir)
x = []
for idx, tensor in enumerate(
zip(input_batches, prediction_batches, label_batches)):
if random.random() > val_args['img_sample_rate']:
continue
pil_input = restore(tensor[0])
pil_output = colorize_mask(tensor[1])
pil_label = colorize_mask(tensor[2])
pil_input.save(os.path.join(to_save_dir, '%d_img.png' % idx))
pil_output.save(os.path.join(to_save_dir, '%d_out.png' % idx))
pil_label.save(os.path.join(to_save_dir, '%d_label.png' % idx))
x.extend([
pil_to_tensor(pil_input.convert('RGB')),
pil_to_tensor(pil_label.convert('RGB')),
pil_to_tensor(pil_output.convert('RGB'))
])
x = torch.stack(x, 0)
x = vutils.make_grid(x, nrow=3, padding=5)
writer.add_image(snapshot_name, x)
print '--------------------------------------------------------'
print '[val loss %.4f], [mean iu %.4f]' % (val_loss, mean_iu)
print '[best val loss %.4f], [mean iu %.4f], [epoch %d]' % (
train_record['best_val_loss'], train_record['corr_mean_iu'],
train_record['corr_epoch'])
print '--------------------------------------------------------'
net.train()
criterion.cuda()