def validate(self, epoch):
start = time.time()
self.net.eval()
val_batch_size = min(ARGS['batch_size'], len(self.val_dataset))
val_dataloader = DataLoader(self.val_dataset,
batch_size=val_batch_size)
epoch_loss = 0.
for batch_index, items in enumerate(val_dataloader):
images, labels, edges = items['image'], items['label'], items[
'edge']
# print('label min:', labels[0].min(), ' max: ', labels[0].max())
# print('edge min:', labels[0].min(), ' max: ', labels[0].max())
if ARGS['gpu']:
labels = labels.cuda()
images = images.cuda()
edges = edges.cuda()
print('image shape:', images.size())
with torch.no_grad():
outputs_edge, outputs = self.net(images)
loss_edge = lovasz_softmax(outputs_edge,
edges) # Lovasz-Softmax loss
loss_seg = lovasz_softmax(outputs, labels) #
loss = ARGS['combine_alpha'] * loss_seg + (
1 - ARGS['combine_alpha']) * loss_edge
pred = torch.max(outputs, dim=1)[1]
iou = torch.sum(pred & labels) / (torch.sum(pred | labels) + 1e-6)
print(
'Validating Epoch: {epoch} [{val_samples}/{total_samples}]\tLoss: {:0.4f}\tIoU: {:0.4f}'
.format(loss.item(),
iou.item(),
epoch=epoch,
val_samples=batch_index * val_batch_size,
total_samples=len(val_dataloader.dataset)))
epoch_loss += loss
# update training loss for each iteration
# self.writer.add_scalar('Train/loss', loss.item(), n_iter)
epoch_loss /= len(val_dataloader)
self.writer.add_scalar('Val/loss', epoch_loss, epoch)
finish = time.time()
print('epoch {} training time consumed: {:.2f}s'.format(
epoch, finish - start))
def forward(self, output, target):
logits = F.softmax(output, dim=1)
loss = lovasz_softmax(logits, target, ignore=self.ignore_index)
return loss
def lovasz_softmax(self, probas, labels, classes='present'):
out = F.softmax(probas, dim=1)
loss = L.lovasz_softmax(out, labels)
return loss
def criterion_reg(pred, label):
return nn.BCELoss()(pred, label) + lovasz_softmax(pred, label, classes=[1])
def train(self, epoch):
start = time.time()
self.net.train()
train_dataloader = DataLoader(self.train_dataset,
batch_size=ARGS['batch_size'],
shuffle=False)
epoch_loss = 0.
for batch_index, items in enumerate(train_dataloader):
images, labels, edges = items['image'], items['label'], items[
'edge']
images = images.float()
labels = labels.long()
edges = edges.long()
if ARGS['gpu']:
labels = labels.cuda()
images = images.cuda()
edges = edges.cuda()
self.optimizer.zero_grad()
outputs_edge, outputs = self.net(images)
# print('output edge min:', outputs_edge[0, 1].min(), ' max: ', outputs_edge[0, 1].max())
# plt.imshow(outputs_edge[0, 1].detach().cpu().numpy() * 255, cmap='gray')
# plt.show()
loss_edge = lovasz_softmax(outputs_edge,
edges) # Lovasz-Softmax loss
loss_seg = lovasz_softmax(outputs, labels) #
loss = ARGS['combine_alpha'] * loss_seg + (
1 - ARGS['combine_alpha']) * loss_edge
loss.backward()
self.optimizer.step()
self.lr_scheduler.step()
n_iter = (epoch - 1) * len(train_dataloader) + batch_index + 1
pred = torch.max(outputs, dim=1)[1]
iou = torch.sum(pred & labels) / (torch.sum(pred | labels) + 1e-6)
# print('edge min:', edges.min(), ' max: ', edges.max())
# print('output edge min:', outputs_edge.min(), ' max: ', outputs_edge.max())
print(
'Training Epoch: {epoch} [{trained_samples}/{total_samples}]\tL_edge: {:0.4f}\tL_seg: {:0.4f}\tL_all: {:0.4f}\tIoU: {:0.4f}\tLR: {:0.4f}'
.format(loss_edge.item(),
loss_seg.item(),
loss.item(),
iou.item(),
self.optimizer.param_groups[0]['lr'],
epoch=epoch,
trained_samples=batch_index * ARGS['batch_size'],
total_samples=len(train_dataloader.dataset)))
epoch_loss += loss.item()
# update training loss for each iteration
# self.writer.add_scalar('Train/loss', loss.item(), n_iter)
for name, param in self.net.named_parameters():
layer, attr = os.path.splitext(name)
attr = attr[1:]
self.writer.add_histogram("{}/{}".format(layer, attr), param,
epoch)
epoch_loss /= len(train_dataloader)
self.writer.add_scalar('Train/loss', epoch_loss, epoch)
finish = time.time()
print('epoch {} training time consumed: {:.2f}s'.format(
epoch, finish - start))
