File size: 5,581 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from __future__ import annotations
from dataclasses import dataclass, field
from functools import cached_property
from typing import (
Any,
Callable,
Dict,
Generic,
Iterable,
List,
Mapping,
Optional,
Tuple,
TypeVar,
)
import torch
from ...diffusion.corruption.d3pm_corruption import D3PMCorruption
from ...diffusion.corruption.sde_lib import SDE, Corruption
from ...diffusion.data.batched_data import BatchedData
R = TypeVar("R")
Diffusable = TypeVar("Diffusable", bound=BatchedData)
def _first(s: Iterable):
return next(iter(s))
@dataclass
class MultiCorruptionConfig:
discrete_corruptions: dict[str, Any] = field(default_factory=dict)
sdes: dict[str, Any] = field(default_factory=dict)
class MultiCorruption(Generic[Diffusable]):
"""Wraps multiple `Corruption` instances to operate on different fields of a State
In the forward process, each field of State is corrupted independently.
In the reverse process, a single score model takes in the entire State and
uses it to estimate the score with respect to each field of the State.
"""
def _get_batch_indices(self, batch: Diffusable) -> Dict[str, torch.Tensor]:
return {k: batch.get_batch_idx(k) for k in self.corrupted_fields}
def __init__(
self,
sdes: Optional[Mapping[str, SDE]] = None,
discrete_corruptions: Optional[Mapping[str, D3PMCorruption]] = None,
):
"""
Args:
sdes: mapping from fields of batch to SDE corruption processes
discrete_corruptions: mapping from fields of batch to discrete corruption processes
"""
if sdes is None:
sdes = {}
if discrete_corruptions is None:
discrete_corruptions = {}
assert (
len(sdes) + len(discrete_corruptions) > 0
), "Must have at least one corruption process."
self._sdes = sdes
self._discrete_corruptions = discrete_corruptions
assert (
set(self._sdes.keys()).intersection(set(self._discrete_corruptions.keys())) == set()
), "SDEs and corruptions have overlapping keys."
self._corruptions: Dict[str, Corruption] = {**self._sdes, **self._discrete_corruptions}
# Make the dict sorted by key (to prevent mismatching checkpoints):
self._corruptions = {k: self._corruptions[k] for k in sorted(self._corruptions.keys())}
# All SDEs must have the same T
T_vals = [corruption.T for corruption in self.corruptions.values()]
assert len(set(T_vals)) == 1
@property
def sdes(self) -> Mapping[str, SDE]:
return self._sdes
@property
def has_discrete_corruptions(self) -> bool:
return len(self.discrete_corruptions) > 0
@property
def discrete_corruptions(self) -> Mapping[str, Corruption]:
return self._discrete_corruptions
@property
def corruptions(self) -> Mapping[str, Corruption]:
return self._corruptions
@property
def corrupted_fields(self) -> List[str]:
return list(self.corruptions.keys())
@cached_property
def T(self) -> float:
return _first(self.corruptions.values()).T
def sample_marginal(self, batch: Diffusable, t) -> Diffusable:
def fn_getter(corruption: Corruption) -> Callable[..., Tuple[torch.Tensor, torch.Tensor]]:
return corruption.sample_marginal
noisy_data = self._apply_corruption_fn(
fn_getter,
x=batch,
batch_idx=self._get_batch_indices(batch),
broadcast=dict(t=t),
)
noisy_batch = batch.replace(**noisy_data)
return noisy_batch
def sde(
self, batch: Diffusable, t: torch.Tensor
) -> Dict[str, Tuple[torch.Tensor, torch.Tensor]]:
"""Get drift and diffusion for each component of the state"""
assert (
not self.has_discrete_corruptions
), "Cannot call `sde` on a MultiCorruption with non-SDE corruptions"
fns = {k: sde.sde for k, sde in self.sdes.items()}
return apply(
fns=fns,
broadcast={"batch": batch, "t": t},
x=batch,
batch_idx=self._get_batch_indices(batch),
)
def _apply_corruption_fn(
self,
fn_getter: Callable[[Corruption], Callable[..., R]],
x: BatchedData,
batch_idx: Mapping[str, torch.LongTensor],
broadcast: Optional[Dict] = None,
apply_to: Optional[Mapping[str, Corruption]] = None,
**kwargs,
) -> Dict[str, R]:
if apply_to is None:
apply_to = self.corruptions
fns = {field_name: fn_getter(corruption) for field_name, corruption in apply_to.items()}
return apply(
fns=fns,
broadcast={**(broadcast or dict()), "batch": x},
x=x,
batch_idx=batch_idx,
**kwargs,
)
def apply(fns: Dict[str, Callable[..., R]], broadcast, **kwargs) -> Dict[str, R]:
"""Apply different function with different argument values to each field.
fns: dict of the form {field_name: function_to_apply}
broadcast: arguments that are identical for every field_name
kwargs: dict of the form {argument_name: {field_name: argument_value}}
"""
return {
field_name: fn(
**{k: v[field_name] for k, v in kwargs.items() if field_name in v},
**(broadcast or dict()),
)
for field_name, fn in fns.items()
}
|