YOLO loss

yolo_loss.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 qiang.zhou <qiang.zhou@yz-gpu029.hogpu.cc>
# Created on 2018-03-05 17:21

import os.path as osp
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import math

CURR_DIR = osp.dirname(__file__)
sys.path.append(osp.join(CURR_DIR, '..'))
from configuration import MODEL_CONFIG, TRAIN_CONFIG
from utils.misc_utils import box_iou


class YOLOLoss(nn.Module):
    #TODO: Add focal loss
    def __init__(self, hmap_shape=MODEL_CONFIG['hmap_shape'], alpha=None, gamma=None, size_average=True):
        super(YOLOLoss, self).__init__()
        self.hmap_shape = hmap_shape
        self.size_average = size_average
        self.lambda_noobj = MODEL_CONFIG['lambda_noobj']
        self.lambda_obj = MODEL_CONFIG['lambda_obj']
        self.lambda_coord = MODEL_CONFIG['lambda_coord']
        self.lambda_iou = MODEL_CONFIG['lambda_iou']
        self.use_gpu = TRAIN_CONFIG['use_gpu']
        self.sqrt_trick = MODEL_CONFIG['sqrt_trick']

    def forward(self, inputs, targets):
        N = inputs.size(0)
        C = inputs.size(1)
        H = inputs.size(2)
        W = inputs.size(3)
        mask = Variable(torch.zeros([N, C, H, W]), requires_grad=False)
        if self.use_gpu:
            inputs = inputs.type('torch.cuda.DoubleTensor')
            targets = targets.type('torch.cuda.DoubleTensor')
            mask = mask.type('torch.cuda.DoubleTensor')
        else:
            inputs = inputs.type('torch.DoubleTensor')
            targets = targets.type('torch.DoubleTensor')
            mask = mask.type('torch.DoubleTensor')

        """
        for i in range(N):
            obj_loc = np.argmax(targets[i, 0, :].data.cpu().numpy())
            row = obj_loc // H
            col = obj_loc %  W
            delta = inputs[i, 0, row, col] - 1.0
            loss[i, 0, row, col] = self.lambda_obj * (delta * delta)

            # Forward and backward of coordinates
            tx, ty, tw, th = targets[i, 1:, row, col].data
            ox, oy, ow, oh = inputs[i, 1:, row, col].data
            delta = targets[i, 1:, row, col] - inputs[i, 1:, row, col]
            loss[i, 1:, row, col] = self.lambda_coord * torch.mul(delta, delta)
            if self.sqrt_trick:
                iou = box_iou([ox, oy, ow*ow, oh*oh], [tx, ty, tw*tw, th*th])
                iou_loss += self.lambda_iou * (iou - 1.0) * (iou - 1.0)
            else:
                iou = box_iou([ox, oy, ow, oh], [tx, ty, tw, th])
                iou_loss += self.lambda_iou * (iou - 1.0) * (iou - 1.0)
        """
        targets_np = targets.data.cpu().numpy()
        for i in range(N):
            obj_loc = np.argmax(targets_np[i, 0, :])
            row, col = obj_loc // H, obj_loc % W
            mask[i, 0, :, :] = math.sqrt(self.lambda_noobj)
            mask[i, 0, row, col] = math.sqrt(self.lambda_obj)
            mask[i, 1:, row, col] = math.sqrt(self.lambda_coord)

        loss = nn.MSELoss(size_average=self.size_average)(inputs*mask, targets*mask)

        #if self.size_average is True:
            #loss = loss / N
            #iou_loss = iou_loss / N
        return loss


if __name__ == "__main__":
    a = YOLOLoss()
    x1 = np.zeros([1, 5, 3, 3])
    x2 = np.zeros([1, 5, 3, 3])
    x1[0, :, 1, 1] = 1, 0.5, 0.5, 1, 1
    x2[0, :, 1, 1] = 1, 0.6, 0.5, 0, 1
    x1 = Variable(torch.from_numpy(x1))
    x2 = Variable(torch.from_numpy(x2))

    loss = a.forward(x2, x1)
    print (loss)