File size: 2,859 Bytes
f15d29e | 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 83 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from typing import Optional, Tuple, Union
import torch
from ...diffusion.corruption.sde_lib import SDE, VESDE, VPSDE
from ...diffusion.data.batched_data import BatchedData
B = Optional[torch.LongTensor]
def wrap_at_boundary(x: torch.Tensor, wrapping_boundary: float) -> torch.Tensor:
"""Wrap x at the boundary given by wrapping_boundary.
Args:
x: tensor of shape (batch_size, dim)
wrapping_boundary: float): wrap at [0, wrapping_boundary] in all dimensions.
Returns:
wrapped_x: tensor of shape (batch_size, dim)
"""
return torch.remainder(
x, wrapping_boundary
) # remainder is the same as mod, but works with negative numbers.
class WrappedSDEMixin:
def sample_marginal(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: torch.LongTensor = None,
batch: Optional[BatchedData] = None,
) -> torch.Tensor:
_super = super()
assert (
isinstance(self, SDE)
and hasattr(_super, "sample_marginal")
and hasattr(self, "wrapping_boundary")
)
if (x > self.wrapping_boundary).any() or (x < 0).any():
# Values outside the wrapping boundary are valid in principle, but could point to an issue in the data preprocessing,
# as typically we assume that the input data is inside the wrapping boundary (e.g., angles between 0 and 2*pi).
print("Warning: Wrapped SDE has received input outside of the wrapping boundary.")
noisy_x = _super.sample_marginal(x=x, t=t, batch_idx=batch_idx, batch=batch)
return self.wrap(noisy_x)
def prior_sampling(
self,
shape: Union[torch.Size, Tuple],
conditioning_data: Optional[BatchedData] = None,
batch_idx: B = None,
) -> torch.Tensor:
_super = super()
assert isinstance(self, SDE) and hasattr(_super, "prior_sampling")
return self.wrap(_super.prior_sampling(shape=shape, conditioning_data=conditioning_data))
def wrap(self, x):
assert isinstance(self, SDE) and hasattr(self, "wrapping_boundary")
return wrap_at_boundary(x, self.wrapping_boundary)
class WrappedVESDE(WrappedSDEMixin, VESDE):
def __init__(
self,
wrapping_boundary: float = 1.0,
sigma_min: float = 0.01,
sigma_max: float = 50.0,
):
super().__init__(sigma_min=sigma_min, sigma_max=sigma_max)
self.wrapping_boundary = wrapping_boundary
class WrappedVPSDE(WrappedSDEMixin, VPSDE):
def __init__(
self,
wrapping_boundary: float = 1.0,
beta_min: float = 0.1,
beta_max: float = 20,
):
super().__init__(beta_min=beta_min, beta_max=beta_max)
self.wrapping_boundary = wrapping_boundary
|