| 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.") |
|
|
| |
| 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), |
| ) |
|
|
