File size: 1,540 Bytes
5e1cd25 |
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 57 58 59 60 61 62 |
#!/usr/bin/python3
# -*- coding: utf-8 -*-
from typing import List, Tuple
import torch
import torch.nn as nn
class DiceLoss(nn.Module):
def __init__(self,
reduction: str = "mean",
eps: float = 1e-6,
):
super(DiceLoss, self).__init__()
self.reduction = reduction
self.eps = eps
if reduction not in ("sum", "mean"):
raise AssertionError(f"param reduction must be sum or mean.")
def forward(self, inputs: torch.Tensor, targets: torch.Tensor):
"""
:param inputs: torch.Tensor, shape: [b, t, 1]. vad prob, after sigmoid activation.
:param targets: shape as `inputs`.
:return:
"""
inputs_ = torch.squeeze(inputs, dim=-1)
targets_ = torch.squeeze(targets, dim=-1)
# shape: [b, t]
intersection = (inputs_ * targets_).sum(dim=-1)
union = (inputs_ + targets_).sum(dim=-1)
# shape: [b,]
dice = (2. * intersection + self.eps) / (union + self.eps)
# shape: [b,]
loss = 1. - dice
# shape: [b,]
if self.reduction == "mean":
loss = torch.mean(loss)
elif self.reduction == "sum":
loss = torch.sum(loss)
else:
raise AssertionError
return loss
def main():
inputs = torch.zeros(size=(1, 198, 1), dtype=torch.float32)
loss_fn = DiceLoss()
loss = loss_fn.forward(inputs, inputs)
print(loss)
return
if __name__ == "__main__":
main()
|