Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,467 Bytes
0fd2f06 |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from abc import ABC, abstractmethod
import torch
class DenoisingLoss(ABC):
@abstractmethod
def __call__(
self, x: torch.Tensor, x_pred: torch.Tensor,
noise: torch.Tensor, noise_pred: torch.Tensor,
alphas_cumprod: torch.Tensor,
timestep: torch.Tensor,
**kwargs
) -> torch.Tensor:
"""
Base class for denoising loss.
Input:
- x: the clean data with shape [B, F, C, H, W]
- x_pred: the predicted clean data with shape [B, F, C, H, W]
- noise: the noise with shape [B, F, C, H, W]
- noise_pred: the predicted noise with shape [B, F, C, H, W]
- alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T]
- timestep: the current timestep with shape [B, F]
"""
pass
class X0PredLoss(DenoisingLoss):
def __call__(
self, x: torch.Tensor, x_pred: torch.Tensor,
noise: torch.Tensor, noise_pred: torch.Tensor,
alphas_cumprod: torch.Tensor,
timestep: torch.Tensor,
**kwargs
) -> torch.Tensor:
return torch.mean((x - x_pred) ** 2)
class VPredLoss(DenoisingLoss):
def __call__(
self, x: torch.Tensor, x_pred: torch.Tensor,
noise: torch.Tensor, noise_pred: torch.Tensor,
alphas_cumprod: torch.Tensor,
timestep: torch.Tensor,
**kwargs
) -> torch.Tensor:
weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1))
return torch.mean(weights * (x - x_pred) ** 2)
class NoisePredLoss(DenoisingLoss):
def __call__(
self, x: torch.Tensor, x_pred: torch.Tensor,
noise: torch.Tensor, noise_pred: torch.Tensor,
alphas_cumprod: torch.Tensor,
timestep: torch.Tensor,
**kwargs
) -> torch.Tensor:
return torch.mean((noise - noise_pred) ** 2)
class FlowPredLoss(DenoisingLoss):
def __call__(
self, x: torch.Tensor, x_pred: torch.Tensor,
noise: torch.Tensor, noise_pred: torch.Tensor,
alphas_cumprod: torch.Tensor,
timestep: torch.Tensor,
**kwargs
) -> torch.Tensor:
return torch.mean((kwargs["flow_pred"] - (noise - x)) ** 2)
NAME_TO_CLASS = {
"x0": X0PredLoss,
"v": VPredLoss,
"noise": NoisePredLoss,
"flow": FlowPredLoss
}
def get_denoising_loss(loss_type: str) -> DenoisingLoss:
return NAME_TO_CLASS[loss_type]
|