seeker
86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
import torch
|
|
|
|
from pi05mini.flow import FlowMatching
|
|
|
|
|
|
def _expand_batch_time(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
"""Expand a per-example time tensor over non-batch dimensions."""
|
|
return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))
|
|
|
|
|
|
def test_training_batch_obeys_linear_flow_equations() -> None:
|
|
flow = FlowMatching()
|
|
target = torch.tensor(
|
|
[
|
|
[[1.0, -2.0], [0.5, 3.0]],
|
|
[[-1.0, 0.0], [2.5, -0.5]],
|
|
]
|
|
)
|
|
generator = torch.Generator().manual_seed(123)
|
|
|
|
x_t, t, target_velocity = flow.training_batch(target, generator=generator)
|
|
|
|
assert x_t.shape == target.shape
|
|
assert target_velocity.shape == target.shape
|
|
assert t.shape[0] == target.shape[0]
|
|
assert torch.all((t >= 0.0) & (t <= 1.0))
|
|
|
|
t_view = _expand_batch_time(t, target)
|
|
# v = target - noise, so the sampled noise is exactly recoverable.
|
|
noise = target - target_velocity
|
|
expected_x_t = (1.0 - t_view) * noise + t_view * target
|
|
torch.testing.assert_close(x_t, expected_x_t)
|
|
|
|
|
|
def test_training_batch_is_reproducible_with_generator() -> None:
|
|
flow = FlowMatching()
|
|
target = torch.randn(3, 5, 2)
|
|
|
|
batch_a = flow.training_batch(
|
|
target, generator=torch.Generator().manual_seed(99)
|
|
)
|
|
batch_b = flow.training_batch(
|
|
target, generator=torch.Generator().manual_seed(99)
|
|
)
|
|
|
|
for tensor_a, tensor_b in zip(batch_a, batch_b):
|
|
torch.testing.assert_close(tensor_a, tensor_b)
|
|
|
|
|
|
def test_loss_is_mean_squared_velocity_error() -> None:
|
|
flow = FlowMatching()
|
|
target_velocity = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
|
|
prediction = target_velocity + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
|
|
|
|
# Squared errors are [1, 1, 4, 0], whose mean is 1.5.
|
|
torch.testing.assert_close(flow.loss(prediction, target_velocity), torch.tensor(1.5))
|
|
torch.testing.assert_close(
|
|
flow.loss(target_velocity, target_velocity), torch.tensor(0.0)
|
|
)
|
|
|
|
|
|
def test_euler_sampler_reaches_fixed_target_for_exact_vector_field() -> None:
|
|
flow = FlowMatching()
|
|
target = torch.tensor(
|
|
[
|
|
[[0.25, -0.5], [1.0, 2.0]],
|
|
[[-2.0, 0.75], [0.0, 0.5]],
|
|
]
|
|
)
|
|
|
|
def exact_model(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
|
|
t_view = _expand_batch_time(t, x_t)
|
|
# Straight conditional path to target. The sampler evaluates t < 1.
|
|
return (target.to(x_t) - x_t) / (1.0 - t_view)
|
|
|
|
torch.manual_seed(5)
|
|
sampled = flow.sample(
|
|
exact_model,
|
|
shape=tuple(target.shape),
|
|
device=target.device,
|
|
steps=8,
|
|
)
|
|
|
|
torch.testing.assert_close(sampled, target, atol=1e-5, rtol=1e-5)
|
|
|