File size: 2,576 Bytes
352cafd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import torch.nn.functional as F
from util.util import compute_tensor_iu
def get_new_iou_hook(values, size):
return 'iou/new_iou_%s'%size, values['iou/new_i_%s'%size]/values['iou/new_u_%s'%size]
def get_orig_iou_hook(values):
return 'iou/orig_iou', values['iou/orig_i']/values['iou/orig_u']
def get_iou_gain(values, size):
return 'iou/iou_gain_%s'%size, values['iou/new_iou_%s'%size] - values['iou/orig_iou']
iou_hooks_to_be_used = [
get_orig_iou_hook,
lambda x: get_new_iou_hook(x, '224'), lambda x: get_iou_gain(x, '224'),
]
iou_hooks_final_only = [
get_orig_iou_hook,
lambda x: get_new_iou_hook(x, '224'), lambda x: get_iou_gain(x, '224'),
]
# Compute common loss and metric for generator only
def compute_loss_and_metrics(images, para, detailed=True, need_loss=True, has_lower_res=True):
"""
This part compute loss and metrics for the generator
"""
loss_and_metrics = {}
gt = images['gt']
seg = images['seg']
pred_224 = images['pred_224']
if need_loss:
# Loss weights
ce_weights = para['ce_weight']
l1_weights = para['l1_weight']
l2_weights = para['l2_weight']
# temp holder for losses at different scale
ce_loss = 0
l1_loss = 0
l2_loss = 0
loss = 0
ce_loss = F.binary_cross_entropy_with_logits(images['out_224'], (gt>0.5).float())
l1_loss = F.l1_loss(pred_224, gt)
l2_loss = F.mse_loss(pred_224, gt)
loss_and_metrics['grad_loss'] = F.l1_loss(images['gt_sobel'], images['pred_sobel'])
# Weighted loss for different levels
loss = ce_loss * ce_weights + l1_loss * l1_weights + l2_loss * l2_weights
loss += loss_and_metrics['grad_loss'] * para['grad_weight']
"""
Compute IOU stats
"""
orig_total_i, orig_total_u = compute_tensor_iu(seg>0.5, gt>0.5)
loss_and_metrics['iou/orig_i'] = orig_total_i
loss_and_metrics['iou/orig_u'] = orig_total_u
new_total_i, new_total_u = compute_tensor_iu(pred_224>0.5, gt>0.5)
loss_and_metrics['iou/new_i_224'] = new_total_i
loss_and_metrics['iou/new_u_224'] = new_total_u
"""
All done.
Now gather everything in a dict for logging
"""
if need_loss:
loss_and_metrics['total_loss'] = 0
loss_and_metrics['ce_loss'] = ce_loss
loss_and_metrics['l1_loss'] = l1_loss
loss_and_metrics['l2_loss'] = l2_loss
loss_and_metrics['loss'] = loss
loss_and_metrics['total_loss'] += loss
return loss_and_metrics
|