Как реализовать взвешенную потерю многозадачности для самостоятельного выполнения задач (Кендалл и др., 2018) в pytorch? ⇐ Python

[Anonymous]

В этом исследовании авторы представляют уравнение (уравнение 7) для взвешивания отдельных потерь для различных задач нейронных сетей.
Я хочу реализовать это как функцию в pytorch чтобы я мог использовать для своей модели. На данный момент я попробовал следующее:

Код: Выделить всё

import torch

# function to apply uncertainty weighing on the losses
def apply_uncertainty_weights(sigma, loss):

"""
This function applies uncertainty weights on the given loss.
NOTE: This implementation is based on the study Kendall et al. 2018 (https://arxiv.org/abs/1705.07115)

Arguments:
sigma: A NN learned uncertainty value (initialised as torch.nn.Parameter(torch.zeros(num_tasks))
loss: The calculated losss between the prediction and target

Returns:
weighted_loss: Weighted loss
"""

# apply uncertainty weighthing
# This is the formula in the publication -> weighted_loss = (1 / (2 * sigma**2)) * loss + torch.log(sigma)
# but we can't use it as it won't be numerically stable/differentiable (e.g. when sigma is predicted to be 0)
# instead use the following
sigma = torch.nn.functional.softplus(sigma) + torch.tensor(1e-8) # this makes sure sigma is never exactly 0 or less otherwise the following functions wont work
log_sigma_squared = torch.log(sigma ** 2) # this is log(sigma^2)
precision = (1/2) * torch.exp(-log_sigma_squared) # this is 1/sigma^2
log_sigma = (1/2) * log_sigma_squared # this is log(sigma)
weighted_loss = precision * loss + log_sigma

# return the weighted loss
return weighted_loss

Но, как ни странно, эта реализация дает мне отрицательные значения потерь во время обучения. Что я делаю не так?

Подробнее здесь: https://stackoverflow.com/questions/792 ... 018-in-pyt