Alex Blank
119 lines
4.6 KiB
Python
119 lines
4.6 KiB
Python
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
|
|
|
|
class QuantileLoss(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.register_buffer('q', torch.tensor(config.quantiles))
|
|
|
|
def forward(self, predictions, targets):
|
|
diff = predictions - targets.unsqueeze(-1).expand(-1, -1, -1, predictions.shape[-1])
|
|
ql = (1 - self.q) * F.relu(diff) + self.q * F.relu(-diff)
|
|
losses = ql.view(-1, ql.shape[-1]).mean(0)
|
|
return losses
|
|
|
|
|
|
def qrisk(pred, tgt, quantiles):
|
|
diff = pred - tgt
|
|
ql = (1 - quantiles) * np.clip(diff, 0, float('inf')) + quantiles * np.clip(-diff, 0, float('inf'))
|
|
losses = ql.reshape(-1, ql.shape[-1])
|
|
normalizer = np.abs(tgt).mean()
|
|
risk = 2 * losses / normalizer
|
|
return risk.mean(0)
|
|
|
|
|
|
def weighted_bce_loss_fn():
|
|
bce = nn.BCEWithLogitsLoss(reduction="none") # Use reduction="none" to get per-element loss
|
|
use_class_weights = False
|
|
time_weight_factor = 1
|
|
|
|
def _loss_fn(y_pred, y_true):
|
|
loss_list = [] # To accumulate per-sample loss values
|
|
for i in range(len(y_pred)):
|
|
current_pred = y_pred[i]
|
|
current_true = y_true[i]
|
|
loss_raw = bce(current_pred, current_true)
|
|
|
|
if use_class_weights:
|
|
# Calculate class weights using torch operations
|
|
try:
|
|
factor = torch.max(current_true) / torch.mean(current_true) / 2
|
|
non_zero_indices = current_true != 0
|
|
class_weights = torch.ones_like(current_pred)
|
|
class_weights[non_zero_indices] = factor
|
|
except Exception:
|
|
class_weights = torch.ones_like(current_pred)
|
|
else:
|
|
class_weights = torch.ones_like(current_pred)
|
|
|
|
# Calculate time weights using torch operations
|
|
try:
|
|
# Find indices where the event occurs
|
|
event_indices = torch.where(current_true == 1)[0]
|
|
if len(event_indices) > 0:
|
|
first_event_index = event_indices[0].item()
|
|
time_weights = torch.ones_like(current_pred)
|
|
# Weight the loss before the first event higher if needed
|
|
time_weights[:first_event_index] = time_weight_factor
|
|
else:
|
|
time_weights = torch.ones_like(current_pred)
|
|
except Exception:
|
|
time_weights = torch.ones_like(current_pred)
|
|
|
|
# Apply the weights to the raw loss
|
|
loss = loss_raw * class_weights * time_weights
|
|
# Optionally, take the mean over the time dimension
|
|
loss_list.append(loss.mean())
|
|
|
|
# Aggregate the loss for the entire batch
|
|
return torch.stack(loss_list).mean()
|
|
|
|
return _loss_fn
|
|
|
|
|
|
def weighted_mse_prob_loss_fn():
|
|
mse = nn.MSELoss(reduce=False)
|
|
after_event_factor = 1
|
|
use_weighted = False
|
|
|
|
def _loss_fn(y_pred, y_true):
|
|
nonlocal mse, after_event_factor
|
|
loss_list = list()
|
|
for i in range(len(y_pred)):
|
|
current_pred = y_pred[i]
|
|
current_true = y_true[i]
|
|
loss_raw = mse(current_pred, current_true)
|
|
# calculate factor to use for weighting, only works with probability values
|
|
if use_weighted:
|
|
try:
|
|
factor = torch.max(current_true) / torch.mean(current_true) / 2
|
|
non_zero_indices = current_true != 0
|
|
class_weights = np.ones(len(current_true))
|
|
class_weights[non_zero_indices] = factor
|
|
except:
|
|
# if there are no non-zero indices, set all weights to 1
|
|
class_weights = np.ones(len(current_true))
|
|
else:
|
|
# if we don't want to weight the loss, set all weights to 1
|
|
class_weights = np.ones(len(current_true))
|
|
|
|
# calculate time weight, here, after the event, the loss should be higher
|
|
try:
|
|
last_event_index = np.where(current_true != 0)[-1]
|
|
time_weights = np.zeros(len(current_true))
|
|
# here, we want to weight the loss after the 'fertility curve' higher, so that we don't get the signal too late
|
|
time_weights[last_event_index:] = after_event_factor
|
|
except:
|
|
# if there are no non-zero indices, set all weights to 1
|
|
time_weights = np.ones(len(current_true))
|
|
|
|
loss = loss_raw * torch.tensor(class_weights) * torch.tensor(time_weights)
|
|
loss_list.append(loss.mean())
|
|
|
|
return torch.stack(loss_list).mean()
|
|
|
|
return _loss_fn
|