File size: 1,859 Bytes
864c14f |
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 54 55 56 |
import torch
from torch import nn
class Loss_VAE(nn.Module):
def __init__(self):
super().__init__()
self.mse = nn.MSELoss(reduction='sum')
def forward(self, recon_x, x, mu, log_var):
mse = self.mse(recon_x, x)
kld = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
loss = mse + kld
return loss
def DiceScore(
y_pred: torch.Tensor,
y: torch.Tensor,
include_background: bool = True,
) -> torch.Tensor:
"""Computes Dice score metric from full size Tensor and collects average.
Args:
y_pred: input data to compute, typical segmentation model output.
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
should be binarized.
y: ground truth to compute mean dice metric. It must be one-hot format and first dim is batch.
The values should be binarized.
include_background: whether to skip Dice computation on the first channel of
the predicted output. Defaults to True.
Returns:
Dice scores per batch and per class, (shape [batch_size, num_classes]).
Raises:
ValueError: when `y_pred` and `y` have different shapes.
"""
y = y.float()
y_pred = y_pred.float()
if y.shape != y_pred.shape:
raise ValueError("y_pred and y should have same shapes.")
# reducing only spatial dimensions (not batch nor channels)
n_len = len(y_pred.shape)
reduce_axis = list(range(2, n_len))
intersection = torch.sum(y * y_pred, dim=reduce_axis)
y_o = torch.sum(y, reduce_axis)
y_pred_o = torch.sum(y_pred, dim=reduce_axis)
denominator = y_o + y_pred_o
return torch.where(
denominator > 0,
(2.0 * intersection) / denominator,
torch.tensor(float("1"), device=y_o.device),
)
|