# -*- coding: utf-8 -*-
# ------------------------
# written by Songjian Chen
# 2018-10
# ------------------------
import torch.nn as nn
import torch
class Compose:
"""Compose several loss functions
"""
def __init__(self, loss_functions):
self.loss_functions = loss_functions
def __call__(self, output, ground_truth):
loss = 0.0
for coefficient, loss_func in self.loss_functions:
loss += coefficient * loss_func(output, ground_truth)
return loss
[docs]class AVGLoss:
""" The loss function used in train part, torch.nn.MSELoss with reduction='mean', which means the sum of the density map will be divided by
the number of pixels in the density map. It can be described as:
.. math::
\ell(\\theta) = \\frac{1}{2N*B}\sum_{i=1}^N\parallel{Z(X_i;\\theta) - Z_i^GT}\parallel_2^2
"""
def __init__(self):
super(AVGLoss).__init__()
[docs] def __call__(self, output, ground_truth):
"""
Args:
output (torch.tensor): the output generated by model
ground_truth (torch.tensor): the ground truth compared with output
Return:
float
"""
batch = len(output)
loss_function = nn.MSELoss(reduction='mean')
loss = loss_function(output.squeeze(), ground_truth.squeeze())
return loss / (2 * batch)
[docs]class SUMLoss:
""" The loss function used in train part, torch.nn.MSELoss with reduction='sum', which means the density map will be summed.
It can be described as:
.. math::
\ell(\\theta) = \\frac{1}{2*B}\sum_{i=1}^N\parallel{Z(X_i;\\theta) - Z_i^GT}\parallel_2^2
"""
def __init__(self):
super(SUMLoss).__init__()
[docs] def __call__(self, output, ground_truth):
"""
Args:
output (torch.tensor): the output generated by model
ground_truth (torch.tensor): the ground truth compared with output
Return:
float
"""
batch = len(output)
loss_function = nn.MSELoss(reduction='sum')
loss = loss_function(output.squeeze(), ground_truth.squeeze())
return loss / (2 * batch)
[docs]class TestLoss:
"""The loss function used in test part, this loss just get mae and mse with (output-ground_truth) and mae ** 2
"""
def __init__(self):
super(TestLoss).__init__()
[docs] def __call__(self, output, ground_truth):
"""
Args:
output (torch.tensor): the output generated by model
ground_truth (torch.tensor): the ground truth compared with output
Return:
float
"""
mae, mse = [0.0] * 2
batch = len(output)
for i in range(batch):
sum_output = torch.sum(output[i])
sum_gt = torch.sum(ground_truth[i])
mae += abs(sum_output - sum_gt)
mse += (sum_output - sum_gt) * (sum_output - sum_gt)
return mae / batch, mse / batch
[docs]class EnlargeLoss:
"""When you enlarge the density map with scale factor (10, 100 or 1000), get test loss with this function.
Args:
number (int): the scale factor used to enlarge the density map
"""
def __init__(self, number):
super(EnlargeLoss).__init__()
self.number = number
[docs] def __call__(self, output, ground_truth):
"""
Args:
output (torch.tensor): the output generated by model, will be divided by self.number
ground_truth (torch.tensor): the ground truth compared with output
Return:
float
"""
mae, mse = [0.0] * 2
batch = len(output)
for i in range(batch):
sum_output = torch.sum(output[i] / self.number)
sum_gt = torch.sum(ground_truth[i] / self.number)
mae += abs(sum_output - sum_gt)
mse += (sum_output - sum_gt) * (sum_output - sum_gt)
return mae / batch, mse / batch
class RankLoss:
"""batch-rank loss of multi-batch training
Args:
"""
def __init__(self, enlarge_number=1):
super(RankLoss).__init__()
self.enlarge_number = enlarge_number
def __call__(self, output, ground_truth):
"""
Args:
output (torch.tensor): the output generated by model
ground_truth (torch.tensor): the ground truth compared with output
Return:
float
"""
loss = 0.0
batch = len(output)
for i in range(batch - 1):
size = ground_truth.size()[1] * ground_truth.size()[2]
sum_output = abs(torch.sum(output[i]) - torch.sum(output[i + 1])) / self.enlarge_number
sum_gt = abs(torch.sum(ground_truth[i]) - torch.sum(ground_truth[i + 1])) / self.enlarge_number
loss += abs(sum_output - sum_gt) / size
return loss