File size: 2,713 Bytes
7f316fe | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
from abc import ABC, abstractmethod
from torch import Tensor
from flow_matching.path.path_sample import PathSample
class ProbPath(ABC):
r"""Abstract class, representing a probability path.
A probability path transforms the distribution :math:`p(X_0)` into :math:`p(X_1)` over :math:`t=0\rightarrow 1`.
The ``ProbPath`` class is designed to support model training in the flow matching framework. It supports two key functionalities: (1) sampling the conditional probability path and (2) conversion between various training objectives.
Here is a high-level example
.. code-block:: python
# Instantiate a probability path
my_path = ProbPath(...)
for x_0, x_1 in dataset:
# Sets t to a random value in [0,1]
t = torch.rand()
# Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
# Optimizes the model. The loss function varies, depending on model and path.
loss(path_sample, my_model(x_t, t)).backward()
"""
@abstractmethod
def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
r"""Sample from an abstract probability path:
| given :math:`(X_0,X_1) \sim \pi(X_0,X_1)`.
| returns :math:`X_0, X_1, X_t \sim p_t(X_t)`, and a conditional target :math:`Y`, all objects are under ``PathSample``.
Args:
x_0 (Tensor): source data point, shape (batch_size, ...).
x_1 (Tensor): target data point, shape (batch_size, ...).
t (Tensor): times in [0,1], shape (batch_size).
Returns:
PathSample: a conditional sample.
"""
def assert_sample_shape(self, x_0, x_1, t):
# x_0, x_1: list[Tensor] of variable length (ragged)
# t: Tensor of shape (B,)
assert isinstance(x_0, (list, tuple)), f"x_0 must be a list or tuple of ragged Tensors. Got {type(x_0)}"
assert isinstance(x_1, (list, tuple)), f"x_1 must be a list or tuple of ragged Tensors. Got {type(x_1)}"
assert isinstance(t, Tensor), f"t must be a torch.Tensor. Got {type(t)}"
batch_size = len(x_0)
assert len(x_1) == batch_size, f"x_1 must match x_0 batch size ({batch_size}). Got {len(x_1)}"
assert t.ndim == 1, f"The time vector t must have shape [batch_size]. Got {t.shape}."
assert t.shape[0] == batch_size, f"Time t dimension must match the batch size [{batch_size}]. Got {t.shape}"
|