Include clean PRIMA code and evaluator; document metrics, initialization, and MIT contributions
91c1e9e verified | 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 | |