File size: 1,131 Bytes
91c1e9e | 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 | from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.llama.modeling_llama import (
LlamaAttention,
apply_rotary_pos_emb,
math,
)
def calculate_dice_loss(
predictions: torch.Tensor,
ground_truth: torch.Tensor,
mask_count: float,
scale_factor=1000,
eps=1e-6,
):
predictions = predictions.sigmoid()
predictions = predictions.flatten(1, 2)
ground_truth = ground_truth.flatten(1, 2)
intersection = 2 * (predictions / scale_factor * ground_truth).sum(dim=-1)
union = (predictions / scale_factor).sum(dim=-1) + (
ground_truth / scale_factor
).sum(dim=-1)
dice_loss = 1 - (intersection + eps) / (union + eps)
dice_loss = dice_loss.sum() / (mask_count + 1e-8)
return dice_loss
def compute_sigmoid_cross_entropy(
predictions: torch.Tensor, targets: torch.Tensor, mask_count: float
):
loss = F.binary_cross_entropy_with_logits(predictions, targets, reduction="none")
loss = loss.flatten(1, 2).mean(1)
loss = loss.sum() / (mask_count + 1e-8)
return loss
|