Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Alex/add eps #124

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Alex/add eps #124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c710150
exposing eps as a parameter
moinnadeem Jul 4, 2024
0b43559
plumb through epsilon fixing + add test
atong01 Jul 10, 2024
cdbbef6
Reformat with black
atong01 Jul 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions tests/test_conditional_flow_matcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,8 @@ def sample_plan(method, x0, x1, sigma):
# Test both integer and floating sigma
@pytest.mark.parametrize("sigma", [0.0, 5e-4, 0.5, 1.5, 0, 1])
@pytest.mark.parametrize("shape", [[1], [2], [1, 2], [3, 4, 5]])
def test_fm(method, sigma, shape):
@pytest.mark.parametrize("test_eps", [False, True])
def test_fm(method, sigma, shape, test_eps):
batch_size = TEST_BATCH_SIZE

if method in SIGMA_CONDITION.keys() and SIGMA_CONDITION[method](sigma):
Expand All @@ -106,7 +107,14 @@ def test_fm(method, sigma, shape):
x0, x1 = random_samples(shape, batch_size=batch_size)
torch.manual_seed(TEST_SEED)
np.random.seed(TEST_SEED)
t, xt, ut, eps = FM.sample_location_and_conditional_flow(x0, x1, return_noise=True)
eps = None
if test_eps:
eps = torch.randn_like(x0)
t, xt, ut, ret_eps = FM.sample_location_and_conditional_flow(
x0, x1, return_noise=True, eps=eps
)
if test_eps:
assert torch.allclose(ret_eps, eps)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this assert should be done at the end of the file. eps is supposed to be the same as ret_eps and it is in the same spirit as the tests over t_given_init

_ = FM.compute_lambda(t)

if method in ["sb_cfm", "exact_ot_cfm"]:
Expand All @@ -115,13 +123,14 @@ def test_fm(method, sigma, shape):
x0, x1 = sample_plan(method, x0, x1, sigma)

torch.manual_seed(TEST_SEED)
if test_eps:
# compute to get same t seed
eps = torch.randn_like(x0)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

where is this one used?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

okay I get it. Why don't you give eps instead of ret_eps in line 132 and drop the comment? esp and ret_eps are supposed the be the same.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we want to test that they are the same? Make sure we didn't mess anything up.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ahh no we can't do that. It messes up the seeds for later inits.

t_given_init = torch.rand(batch_size)
t_given = t_given_init.reshape(-1, *([1] * (x0.dim() - 1)))
sigma_pad = pad_t_like_x(sigma, x0)
epsilon = torch.randn_like(x0)
computed_xt, computed_ut = compute_xt_ut(method, x0, x1, t_given, sigma_pad, epsilon)
computed_xt, computed_ut = compute_xt_ut(method, x0, x1, t_given, sigma_pad, ret_eps)

assert torch.all(ut.eq(computed_ut))
assert torch.all(xt.eq(computed_xt))
assert torch.all(eps.eq(epsilon))
assert torch.allclose(xt, computed_xt)
kilianFatras marked this conversation as resolved.
Show resolved Hide resolved
assert any(t_given_init == t)
52 changes: 32 additions & 20 deletions torchcfm/conditional_flow_matching.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,10 +152,7 @@ def compute_conditional_flow(self, x0, x1, t, xt):
del t, xt
return x1 - x0

def sample_noise_like(self, x):
return torch.randn_like(x)

def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False):
def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False, eps=None):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you run some notebooks to ensure the behaviour is still correct please? thx.

"""
Compute the sample xt (drawn from N(t * x1 + (1 - t) * x0, sigma))
and the conditional vector field ut(x1|x0) = x1 - x0, see Eq.(15) [1].
Expand All @@ -169,8 +166,10 @@ def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=Fals
(optionally) t : Tensor, shape (bs)
represents the time levels
if None, drawn from uniform [0,1]
return_noise : bool
(optionally) return_noise : bool
return the noise sample epsilon
(optionally) eps: Tensor, shape (bs, *dim)
use a fixed noise vector epsilon


Returns
Expand All @@ -189,7 +188,8 @@ def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=Fals
t = torch.rand(x0.shape[0]).type_as(x0)
assert len(t) == x0.shape[0], "t has to have batch size dimension"

eps = self.sample_noise_like(x0)
if eps is None:
eps = torch.randn_like(x0)
xt = self.sample_xt(x0, x1, t, eps)
ut = self.compute_conditional_flow(x0, x1, t, xt)
if return_noise:
Expand Down Expand Up @@ -234,7 +234,7 @@ def __init__(self, sigma: Union[float, int] = 0.0):
super().__init__(sigma)
self.ot_sampler = OTPlanSampler(method="exact")

def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False):
def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False, eps=None):
r"""
Compute the sample xt (drawn from N(t * x1 + (1 - t) * x0, sigma))
and the conditional vector field ut(x1|x0) = x1 - x0, see Eq.(15) [1]
Expand All @@ -249,8 +249,10 @@ def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=Fals
(optionally) t : Tensor, shape (bs)
represents the time levels
if None, drawn from uniform [0,1]
return_noise : bool
(optionally) return_noise : bool
return the noise sample epsilon
(optionally) eps: Tensor, shape (bs, *dim)
use a fixed noise vector epsilon

Returns
-------
Expand All @@ -265,10 +267,10 @@ def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=Fals
[1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.
"""
x0, x1 = self.ot_sampler.sample_plan(x0, x1)
return super().sample_location_and_conditional_flow(x0, x1, t, return_noise)
return super().sample_location_and_conditional_flow(x0, x1, t, return_noise, eps)

def guided_sample_location_and_conditional_flow(
self, x0, x1, y0=None, y1=None, t=None, return_noise=False
self, x0, x1, y0=None, y1=None, t=None, return_noise=False, eps=None
):
r"""
Compute the sample xt (drawn from N(t * x1 + (1 - t) * x0, sigma))
Expand All @@ -288,8 +290,10 @@ def guided_sample_location_and_conditional_flow(
(optionally) t : Tensor, shape (bs)
represents the time levels
if None, drawn from uniform [0,1]
return_noise : bool
(optionally) return_noise : bool
return the noise sample epsilon
(optionally) eps: Tensor, shape (bs, *dim)
use a fixed noise vector epsilon

Returns
-------
Expand All @@ -305,10 +309,12 @@ def guided_sample_location_and_conditional_flow(
"""
x0, x1, y0, y1 = self.ot_sampler.sample_plan_with_labels(x0, x1, y0, y1)
if return_noise:
t, xt, ut, eps = super().sample_location_and_conditional_flow(x0, x1, t, return_noise)
t, xt, ut, eps = super().sample_location_and_conditional_flow(
x0, x1, t, return_noise, eps
)
return t, xt, ut, y0, y1, eps
else:
t, xt, ut = super().sample_location_and_conditional_flow(x0, x1, t, return_noise)
t, xt, ut = super().sample_location_and_conditional_flow(x0, x1, t, return_noise, eps)
return t, xt, ut, y0, y1


Expand Down Expand Up @@ -468,7 +474,7 @@ def compute_conditional_flow(self, x0, x1, t, xt):
ut = sigma_t_prime_over_sigma_t * (xt - mu_t) + x1 - x0
return ut

def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False):
def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False, eps=None):
"""
Compute the sample xt (drawn from N(t * x1 + (1 - t) * x0, sqrt(t * (1 - t))*sigma^2 ))
and the conditional vector field ut(x1|x0) = (1 - 2 * t) / (2 * t * (1 - t)) * (xt - mu_t) + x1 - x0,
Expand All @@ -483,8 +489,10 @@ def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=Fals
(optionally) t : Tensor, shape (bs)
represents the time levels
if None, drawn from uniform [0,1]
return_noise: bool
(optionally) return_noise: bool
return the noise sample epsilon
(optionally) eps: Tensor, shape (bs, *dim)
use a fixed noise vector epsilon


Returns
Expand All @@ -500,10 +508,10 @@ def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=Fals
[1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.
"""
x0, x1 = self.ot_sampler.sample_plan(x0, x1)
return super().sample_location_and_conditional_flow(x0, x1, t, return_noise)
return super().sample_location_and_conditional_flow(x0, x1, t, return_noise, eps)

def guided_sample_location_and_conditional_flow(
self, x0, x1, y0=None, y1=None, t=None, return_noise=False
self, x0, x1, y0=None, y1=None, t=None, return_noise=False, eps=None
):
r"""
Compute the sample xt (drawn from N(t * x1 + (1 - t) * x0, sigma))
Expand All @@ -523,8 +531,10 @@ def guided_sample_location_and_conditional_flow(
(optionally) t : Tensor, shape (bs)
represents the time levels
if None, drawn from uniform [0,1]
return_noise : bool
(optionally) return_noise : bool
return the noise sample epsilon
(optionally) eps: Tensor, shape (bs, *dim)
use a fixed noise vector epsilon

Returns
-------
Expand All @@ -540,10 +550,12 @@ def guided_sample_location_and_conditional_flow(
"""
x0, x1, y0, y1 = self.ot_sampler.sample_plan_with_labels(x0, x1, y0, y1)
if return_noise:
t, xt, ut, eps = super().sample_location_and_conditional_flow(x0, x1, t, return_noise)
t, xt, ut, eps = super().sample_location_and_conditional_flow(
x0, x1, t, return_noise, eps
)
return t, xt, ut, y0, y1, eps
else:
t, xt, ut = super().sample_location_and_conditional_flow(x0, x1, t, return_noise)
t, xt, ut = super().sample_location_and_conditional_flow(x0, x1, t, return_noise, eps)
return t, xt, ut, y0, y1


Expand Down
Loading