Skip to content

Adding a few safety features #19

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

Open
stevet40 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Adding a few safety features #19

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
75 changes: 41 additions & 34 deletions flowfusion/diffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,9 @@ class ScoreModel(torch.nn.Module):
Internal variable for tracking conditional inputs.
no_sigma : bool
If `True`, `model` is assumed to return score(x, t, conditional).
This is the default behaviour for an `SDE` of class `VPSDE`.
If `False`, `model` is assumed to return score(x, t, conditional) * sigma(t).
This is the default behaviour for an `SDE` of class `VESDE`.
prob : bool
Internal variable to track whether the trace of the Jacobian is included
in the forward call (automatically set/reset when calling `solve_odes_forward`).
Expand All @@ -160,7 +162,6 @@ def __init__(
model=None,
sde=None,
conditional=None,
no_sigma=False,
hutchinson=False,
hutchpp=False,
hpp_rank = 1,
Expand All @@ -177,9 +178,6 @@ def __init__(
Stochastic differential equation. Usually a `VPSDE`, `VESDE` or `SUBVPSDE`.
conditional : torch.Tensor, optional
Initial value of conditioning variable (can be updated)
no_sigma : bool, optional
If `True`, `model` is assumed to return score(x, t, conditional).
If `False`, `model` is assumed to return score(x, t, conditional) * sigma(t).
hutchinson : bool, optional
If `True`, `solve_odes_forward` will be computed using the
Skilling--Hutchinson trace estimator.
Expand All @@ -201,9 +199,11 @@ def __init__(
self.model = model # see above
self.sde = sde
self.conditional = conditional # stores the conditioning variable
self.no_sigma = (
no_sigma # if True, drops the 1/sigma(t) in the score definition
)
# set no_sigma flag automatically now
if isinstance(sde, VESDE):
self.no_sigma = False
else:
self.no_sigma = True # if True, drops the 1/sigma(t) in the score definition
self.prob = False
self.hutch = hutchinson # if True, uses the Hutchinson trace estimator
self.hutchpp = hutchpp # if True, uses the Hutch++ trace estimator
Expand Down Expand Up @@ -570,7 +570,7 @@ def sample_ode_from_base(
atol=1e-4,
rtol=1e-4,
method="dopri5",
options=None,
options={"min_step": 1e-6},
):
"""
Generate samples deterministically by solving ODE backwards in time from t=T to t=0.
Expand Down Expand Up @@ -601,6 +601,11 @@ def sample_ode_from_base(
`torchdiffeq.odeint_adjoint` : ODE solver used when backward pass needed
"""

# check atol/rtol
if atol > 1e-2 or rtol > 1e-2 or atol <= 0 or rtol <= 0:
print("WARNING: suspicious ODE solver tolerance detected!")
print("WARNING: found atol={:.3e} and rtol={:.3e}.".format(atol, rtol))

# base samples (x(t))
if hasattr(self.sde, "sigma_max"):
z = base_samples * self.sde.sigma_max
Expand Down Expand Up @@ -644,10 +649,10 @@ def solve_odes_forward(
self,
x0_samples,
conditional=None,
atol=1e-5,
rtol=1e-5,
atol=1e-4,
rtol=1e-4,
method="dopri5",
options=None,
options={"min_step": 1e-6},
):
"""
This solves the pair of ODEs forward in time to find the base samples, x(t=T),
Expand Down Expand Up @@ -693,6 +698,11 @@ def solve_odes_forward(
`sample_ode_from_base` : Solves in the opposite direction (from t=T to t=0).
"""

# check atol/rtol
if atol > 1e-2 or rtol > 1e-2 or atol <= 0 or rtol <= 0:
print("WARNING: low ODE solver tolerance detected!")
print("WARNING: found atol={:.1e} and rtol={:.1e}.".format(atol, rtol))

# set prob to True
self.prob = True

Expand Down Expand Up @@ -1439,9 +1449,8 @@ def __init__(
shift=None,
scale=None,
method="dopri5",
no_sigma=False,
hutchinson=False,
options=None,
options={"min_step": 1e-6},
):
"""
Parameters
Expand All @@ -1456,10 +1465,6 @@ def __init__(
Parameter scale for inputs/outputs.
method : str, optional
Name of ODE solver. Must be a valid `torchdiffeq` solver name.
no_sigma : bool, optional
If `True`, `model` is assumed to return score(x, t, conditional).
If `False`, `model` is assumed to return score(x, t, conditional) * sigma(t).
For a `VPSDE`, setting `no_sigma=True` is strongly recommended.
hutchinson : bool, optional
If `True`, `log_prob` will be computed using the Skilling--Hutchinson
trace estimator.
Expand All @@ -1470,9 +1475,7 @@ def __init__(

self.model = model
self.sde = sde
self.score_model = ScoreModel(
model=self.model, sde=self.sde, hutchinson=hutchinson, no_sigma=no_sigma
)
self.score_model = ScoreModel(model=self.model, sde=self.sde, hutchinson=hutchinson)
self.register_buffer(
"shift",
(
Expand All @@ -1492,7 +1495,7 @@ def __init__(
self.method = method
self.options = options

def forward(self, base_samples):
def forward(self, base_samples, atol=1e-4, rtol=1e-4):
"""
Generate samples deterministically via an ODE solve.
Applies any rescaling set by `self.shift` and `self.scale`.
Expand All @@ -1506,6 +1509,10 @@ def forward(self, base_samples):
-------
target_samples : torch.Tensor
Rescaled samples from target density, i.e., shift + x(t=0)*scale.
atol : float, optional
Absolute error tolerance for ODE solver.
rtol : float, optional
Relative error tolerance for ODE solver.

See Also
--------
Expand All @@ -1515,8 +1522,8 @@ def forward(self, base_samples):
self.score_model.sample_ode_from_base(
base_samples,
method=self.method,
atol=1e-5,
rtol=1e-5,
atol=atol,
rtol=rtol,
options=self.options,
)[0]
* self.scale
Expand Down Expand Up @@ -1546,7 +1553,7 @@ def sample_sde(self, shape, steps=100):
"""
return self.score_model.sample_sde(shape, steps=100) * self.scale + self.shift

def log_prob(self, x, atol=1e-5, rtol=1e-5):
def log_prob(self, x, atol=1e-4, rtol=1e-4):
"""
Compute log probability of target samples, p[x(t=0)].

Expand Down Expand Up @@ -1622,9 +1629,8 @@ def __init__(
scale=None,
conditional_shift=None,
conditional_scale=None,
no_sigma=False,
method="dopri5",
options=None,
options={"min_step": 1e-6},
):
"""
Parameters
Expand All @@ -1641,9 +1647,6 @@ def __init__(
Parameter shift for conditional inputs.
conditional_scale : torch.Tensor, optional
Parameter scale for conditional inputs.
no_sigma : bool, optional
If `True`, `model` is assumed to return score(x, t, conditional).
If `False`, `model` is assumed to return score(x, t, conditional) * sigma(t).
method : str, optional
Name of ODE solver. Must be a valid `torchdiffeq` solver name.
options : dict, optional
Expand All @@ -1654,7 +1657,7 @@ def __init__(

self.model = model
self.sde = sde
self.score_model = ScoreModel(model=self.model, sde=self.sde, no_sigma=no_sigma)
self.score_model = ScoreModel(model=self.model, sde=self.sde)
self.register_buffer(
"shift",
(
Expand Down Expand Up @@ -1690,7 +1693,7 @@ def __init__(
self.options = options
self.method = method

def forward(self, base_samples, conditional=None):
def forward(self, base_samples, conditional=None, atol=1e-4, rtol=1e-4):
"""
Generate samples deterministically via an ODE solve.
Applies any rescaling set by `self.shift` and `self.scale`.
Expand All @@ -1701,6 +1704,10 @@ def forward(self, base_samples, conditional=None):
Samples from base density, i.e., x(t=T) ~ N(0,1).
conditional : torch.Tensor, optional
Conditional inputs.
atol : float, optional
Absolute error tolerance for ODE solver.
rtol : float, optional
Relative error tolerance for ODE solver.

Returns
-------
Expand All @@ -1714,8 +1721,8 @@ def forward(self, base_samples, conditional=None):
conditional=(conditional - self.conditional_shift)
/ self.conditional_scale,
method=self.method,
atol=1e-5,
rtol=1e-5,
atol=atol,
rtol=rtol,
options=self.options,
)[0]
* self.scale
Expand Down Expand Up @@ -1752,7 +1759,7 @@ def sample_sde(self, shape, conditional=None, steps=100):
+ self.shift
)

def log_prob(self, x, conditional=None, atol=1e-5, rtol=1e-5):
def log_prob(self, x, conditional=None, atol=1e-4, rtol=1e-4):
"""
Compute conditional log probability of target samples, p[x(t=0)|cond].

Expand Down