Spaces:
Running on Zero
Running on Zero
| from dataclasses import dataclass | |
| import torch | |
| from einops import rearrange | |
| from jaxtyping import Float | |
| from lpips import LPIPS | |
| from torch import Tensor | |
| from ..dataset.types import BatchedExample | |
| from ..model.decoder.decoder import DecoderOutput | |
| from ..model.types import Gaussians | |
| from .loss import Loss | |
| class LossLpipsCfg: | |
| weight: float | |
| apply_after_step: int | |
| class LossLpipsCfgWrapper: | |
| lpips: LossLpipsCfg | |
| class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]): | |
| lpips: LPIPS | |
| def __init__(self, cfg: LossLpipsCfgWrapper) -> None: | |
| super().__init__(cfg) | |
| lpips = LPIPS(net="vgg").eval() | |
| lpips.requires_grad_(False) | |
| self.__dict__["lpips"] = lpips | |
| def lpips_on_device(self, device: torch.device) -> LPIPS: | |
| lpips = self.lpips | |
| param = next(iter(lpips.parameters())) | |
| if param.device != device or param.dtype != torch.float32: | |
| lpips.to(device=device, dtype=torch.float32) | |
| return lpips | |
| def forward( | |
| self, | |
| prediction: DecoderOutput, | |
| batch: BatchedExample, | |
| gaussians: Gaussians, | |
| global_step: int, | |
| valid_depth_mask: Tensor | None, | |
| reduction: str = "mean", | |
| ) -> Float[Tensor, "*batch"]: | |
| image = batch["target"]["image"] | |
| # Before the specified step, don't apply the loss. | |
| if global_step < self.cfg.apply_after_step: | |
| return torch.tensor(0, dtype=torch.float32, device=image.device) | |
| prediction_color = prediction.color | |
| if valid_depth_mask is not None and valid_depth_mask.max() > 0.5: | |
| prediction_color = prediction_color.clone() | |
| image = image.clone() | |
| prediction_color[valid_depth_mask] = 0 | |
| image[valid_depth_mask] = 0 | |
| with torch.autocast(device_type=image.device.type, enabled=False): | |
| loss = self.lpips_on_device(image.device).forward( | |
| rearrange(prediction_color.float(), "b v c h w -> (b v) c h w"), | |
| rearrange(image.float(), "b v c h w -> (b v) c h w"), | |
| normalize=True, | |
| ) | |
| batch_size, num_views = image.shape[:2] | |
| sample_loss = loss.reshape(batch_size, num_views, -1).mean(dim=(1, 2)) | |
| if reduction == "mean": | |
| return self.cfg.weight * sample_loss.mean() | |
| if reduction == "none": | |
| return self.cfg.weight * sample_loss | |
| raise ValueError(f"Unsupported reduction={reduction}") | |