def valid(net,
optimizer,
valid_loader,
use_lovasz=False,
save_imgs=False,
fold_num=0):
net.eval()
# keep track of losses
val_ious = []
val_iter_loss = 0.
# no gradients during validation
with torch.no_grad():
for i, data in enumerate(valid_loader):
valid_imgs = data['img'].to(device)
valid_msks = data['msk'].to(device)
valid_msk_bool = data['has_msk'].float().to(device)
# get predictions
msk_vpreds, bool_vpreds = net(valid_imgs)
msk_blend_vpreds = msk_vpreds * bool_vpreds.view(
valid_imgs.size(0), 1, 1, 1)
if save_imgs:
img_grid = vsn.utils.make_grid(valid_imgs, normalize=True)
msk_grid = vsn.utils.make_grid(msk_vpreds)
vsn.utils.save_image(
img_grid,
'../imgs/valid_imgs_fold-{}.png'.format(fold_num))
vsn.utils.save_image(
msk_grid,
'../imgs/valid_msks_fold-{}.png'.format(fold_num))
# calculate loss
if use_lovasz:
vloss = L.lovasz_hinge(msk_vpreds, valid_msks)
#vloss += focal_loss(msk_vpreds, valid_msks)
#vloss -= dice(msk_vpreds.sigmoid(), valid_msks)
else:
vloss = focal_loss(msk_vpreds, valid_msks)
vloss += L.lovasz_hinge(msk_vpreds, valid_msks)
#vloss -= dice(msk_vpreds.sigmoid(), valid_msks)
vloss += args.lambda_bool * bce(bool_vpreds,
valid_msk_bool.view(-1, 1))
vloss += args.lambda_bool * focal_loss(msk_blend_vpreds,
valid_msks)
#vloss += args.lambda_dice * dice(msk_vpreds.sigmoid(), valid_msks)
# get validation stats
val_iter_loss += vloss.item()
val_ious.append(
get_iou_vector(
valid_msks.cpu().numpy()[:, :, 13:114, 13:114],
msk_vpreds.sigmoid().cpu().numpy()[:, :, 13:114, 13:114]))
epoch_vloss = val_iter_loss / (len(valid_loader.dataset) / args.batch_size)
print('Avg Eval Loss: {:.4}, Avg IOU: {:.4}'.format(
epoch_vloss, np.mean(val_ious)))
return epoch_vloss, np.mean(val_ious)
def valid(net, optimizer, valid_loader, mtype='student', use_lovasz=False):
net.eval()
# keep track of losses
val_ious = []
val_iter_loss = 0.
# no gradients during validation
with torch.no_grad():
for i, data in enumerate(valid_loader):
valid_imgs = data['img'].cuda()
valid_msks = data['msk'].cuda()
valid_msk_bool = data['has_msk'].float().cuda()
# get predictions
msk_vpreds, bool_v = net(valid_imgs)
# calculate loss
if use_lovasz:
vloss = L.lovasz_hinge(msk_vpreds, valid_msks)
else:
vloss = focal_loss(msk_vpreds, valid_msks)
vloss += L.lovasz_hinge(msk_vpreds, valid_msks)
vloss += bce(bool_v, valid_msk_bool.view(-1, 1))
# get validation stats
val_iter_loss += vloss.item()
val_ious.append(
get_iou_vector(
valid_msks.cpu().numpy()[:, :, 13:114, 13:114],
msk_vpreds.sigmoid().cpu().numpy()[:, :, 13:114, 13:114]))
epoch_vloss = val_iter_loss / (len(valid_loader.dataset) / args.batch_size)
print('{} Avg Eval Loss: {:.4}, Avg IOU: {:.4}'.format(
mtype, epoch_vloss, np.mean(val_ious)))
return epoch_vloss, np.mean(val_ious)
def forward(self, output_weight, seg_feat, ind, target):
B, C, H, W = seg_feat.size()
weight = _tranpose_and_gather_feat(output_weight, ind)
seg_feat = seg_feat.view([-1, H, W]).unsqueeze(0)
weight = weight.view([-1, C, 3, 3])
pred_seg = F.conv2d(seg_feat, weight, stride=1, padding=1,
groups=B).view([B, -1, H, W])
loss = (L.lovasz_hinge(pred_seg, target, 255) +
L.lovasz_hinge(-pred_seg, 1 - target, -254)) / 2
return loss
def valid(net, optimizer, valid_loader, save_imgs=False, fold_num=0):
net.eval()
# keep track of losses
val_ious = []
val_iter_loss = 0.
# no gradients during validation
with torch.no_grad():
for i, data in enumerate(valid_loader):
valid_imgs = data['img'].to(device)
valid_msks = data['msk']
#valid_msks_half = data['msk_half'].to(device)
#valid_msks_qrtr = data['msk_qrtr'].to(device)
valid_msk_bool = data['has_msk'].to(device)
# get predictions
msk_vpreds = net(valid_imgs)
# calculate loss
# main mask loss
vloss = focal_loss(msk_vpreds[0], valid_msks[0].to(device))
vloss += L.lovasz_hinge(msk_vpreds[0], valid_msks[0].to(device))
# aux mask losses
for j in range(len(valid_msks)):
#print(msk_preds[i+1].size(), msks[i].size())
vloss += 0.1 * focal_loss(
msk_vpreds[j + 1][valid_msk_bool],
valid_msks[j][valid_msk_bool].to(device))
vloss += 0.1 * L.lovasz_hinge(
msk_vpreds[j + 1][valid_msk_bool],
valid_msks[j][valid_msk_bool].to(device))
# binary loss
vloss += 0.05 * bce(msk_vpreds[-1],
valid_msk_bool.float().view(-1, 1))
# get validation stats
val_iter_loss += vloss.item()
val_ious.append(
get_iou_vector(
valid_msks[0].cpu().numpy()[:, :, 13:114, 13:114],
msk_vpreds[0].sigmoid().cpu().numpy()[:, :, 13:114,
13:114]))
epoch_vloss = val_iter_loss / (len(valid_loader.dataset) / args.batch_size)
print('Avg Eval Loss: {:.4}, Avg IOU: {:.4}'.format(
epoch_vloss, np.mean(val_ious)))
return epoch_vloss, np.mean(val_ious)
def train(net,
optimizer,
train_loader,
batch_size,
freeze_bn=False,
use_lovasz=False,
swa=False):
'''
uses the data loader to grab a batch of images
pushes images through network and gathers predictions
updates network weights by evaluating the loss functions
'''
# set network to train mode
net.train(True, args.freeze_bn)
# keep track of our loss
iter_loss = 0.
# loop over the images for the desired amount
for i, data in enumerate(train_loader):
imgs = data['img'].cuda(async=True)
msks = data['msk'].cuda(async=True)
msk_bool = data['has_msk'].float().cuda(async=True)
if args.debug and i == 0:
img_grid = vsn.utils.make_grid(imgs, normalize=True)
msk_grid = vsn.utils.make_grid(msks)
vsn.utils.save_image(img_grid, '../imgs/train_imgs.png')
vsn.utils.save_image(msk_grid, '../imgs/train_msks.png')
# zero gradients from previous run
optimizer.zero_grad()
# get predictions
msk_preds, chck_preds = net(imgs)
# calculate loss
if use_lovasz:
loss = L.lovasz_hinge(msk_preds, msks)
else:
loss = focal_loss(msk_preds, msks)
loss -= dice(msk_preds.sigmoid(), msks)
loss += bce(chck_preds, msk_bool.view(-1, 1))
#loss += args.lambda_dice * dice(msk_preds.sigmoid(), msks)
#calculate gradients
loss.backward()
# update weights
optimizer.step()
# get training stats
iter_loss += loss.item()
# make a cool terminal output
sys.stdout.write('\r')
sys.stdout.write('B: {:>3}/{:<3} | {:.4}'.format(
i + 1, len(train_loader), loss.item()))
epoch_loss = iter_loss / (len(train_loader.dataset) / batch_size)
print('\n' + 'Avg Train Loss: {:.4}'.format(epoch_loss))
return epoch_loss
def train(student,
teacher,
optimizer,
train_loader,
w_t=0.,
e=0,
freeze_bn=False,
use_lovasz=False):
'''
uses the data loader to grab a batch of images
pushes images through network and gathers predictions
updates network weights by evaluating the loss functions
'''
# set network to train mode
student.train(True, freeze_bn)
teacher.train(True, freeze_bn)
# keep track of our loss
iter_loss = 0.
iter_closs = 0.
# loop over the images for the desired amount
for i, data in enumerate(train_loader):
imgs_a = data['img_a'].cuda()
imgs_b = data['img_b'].cuda()
msks = data['msk'].cuda()
labeled_bool = data['is_labeled'].cuda()
has_msk = data['has_msk'].float().cuda()
mask = labeled_bool == 1
if args.debug and i == 0:
#print('{} labeled, {} total'.format(len(imgs_a[mask]), len(imgs_a)))
img_a_grid = vsn.utils.make_grid(imgs_a, normalize=True)
img_b_grid = vsn.utils.make_grid(imgs_b, normalize=True)
msk_grid = vsn.utils.make_grid(msks)
vsn.utils.save_image(img_a_grid, '../imgs/train_mt_imgs_a.png')
vsn.utils.save_image(img_b_grid, '../imgs/train_mt_imgs_b.png')
vsn.utils.save_image(msk_grid, '../imgs/train_msks.png')
# zero gradients from previous run
optimizer.zero_grad()
# get predictions
preds_a, bool_a = student(imgs_a)
# calculate bce loss
if use_lovasz:
loss_s = L.lovasz_hinge(preds_a[mask], msks[mask])
else:
loss_s = focal_loss(preds_a[mask], msks[mask])
loss_s += L.lovasz_hinge(preds_a[mask], msks[mask])
loss_s += bce(bool_a[mask], has_msk[mask].view(-1, 1))
# get the teacher predictions
with torch.no_grad():
preds_b, bool_b = teacher(imgs_b)
loss_c = cl(preds_a, preds_b)
loss_c += cl(bool_a, bool_b)
loss = loss_s + w_t * loss_c
#calculate gradients
loss.backward()
# update weights
optimizer.step()
# get training stats
iter_loss += loss_s.item()
iter_closs += loss_c.item()
# update the teacher weights
grad_step = e * (len(train_loader.dataset) / args.batch_size) + (i + 1)
alpha = min(1. - 1. / (grad_step + 1), 0.99)
update_teacher(teacher, student, alpha=alpha)
# make a cool terminal output
sys.stdout.write('\r')
sys.stdout.write(
'B: {:>3}/{:<3} | loss: {:.4} | step: {} alpha: {:.4}'.format(
i + 1, len(train_loader), loss.item(), int(grad_step), alpha))
epoch_loss = iter_loss / (len(train_loader.dataset) / args.batch_size)
epoch_closs = iter_closs / (len(train_loader.dataset) / args.batch_size)
print('\n' + 'Avg Train Loss: {:.4}, Avg Consist. Loss: {:.4}'.format(
epoch_loss, epoch_closs))
return epoch_loss
def train(net, optimizer, train_loader, freeze_bn=False, swa=False):
'''
uses the data loader to grab a batch of images
pushes images through network and gathers predictions
updates network weights by evaluating the loss functions
'''
# set network to train mode
#if args.gpu == 99:
net.train()
#else:
# net.train(mode=True, freeze_bn=args.freeze_bn)
# keep track of our loss
iter_loss = 0.
# loop over the images for the desired amount
for i, data in enumerate(train_loader):
imgs = data['img'].to(device)
msks = data['msk']
msk_bool = data['has_msk'].to(device)
if args.debug and i == 0:
img_grid = vsn.utils.make_grid(imgs, normalize=True)
msk_grid = vsn.utils.make_grid(msks)
#msk_hgrid = vsn.utils.make_grid(msks_half)
#msk_qgrid = vsn.utils.make_grid(msks_qrtr)
vsn.utils.save_image(img_grid, '../imgs/train_imgs.png')
vsn.utils.save_image(msk_grid, '../imgs/train_msks.png')
#vsn.utils.save_image(msk_hgrid, '../imgs/train_msks_half.png')
#vsn.utils.save_image(msk_qgrid, '../imgs/train_msks_qrtr.png')
# get predictions
msk_preds = net(imgs)
#print(len(msk_preds), len(msks))
# calculate loss
# main mask loss
loss = focal_loss(msk_preds[0], msks[0].to(device))
loss += L.lovasz_hinge(msk_preds[0], msks[0].to(device))
# aux mask losses
for j in range(len(msks)):
#print(msk_preds[j+1].size(), msks[j].size())
loss += 0.1 * focal_loss(msk_preds[j + 1][msk_bool],
msks[j][msk_bool].to(device))
loss += 0.1 * L.lovasz_hinge(msk_preds[j + 1][msk_bool],
msks[j][msk_bool].to(device))
# binary loss
loss += 0.05 * bce(msk_preds[-1], msk_bool.float().view(-1, 1))
# zero gradients from previous run
optimizer.zero_grad()
#calculate gradients
loss.backward()
# update weights
optimizer.step()
# get training stats
iter_loss += loss.item()
# make a cool terminal output
sys.stdout.write('\r')
sys.stdout.write('B: {:>3}/{:<3} | {:.4}'.format(
i + 1, len(train_loader), loss.item()))
epoch_loss = iter_loss / (len(train_loader.dataset) / args.batch_size)
print('\n' + 'Avg Train Loss: {:.4}'.format(epoch_loss))
return epoch_loss