Spaces:
Running
on
Zero
Running
on
Zero
Commit
•
eeb2ff8
1
Parent(s):
b5222e2
Upload 5 files
Browse files- transport/__init__.py +76 -0
- transport/integrators.py +127 -0
- transport/path.py +206 -0
- transport/transport.py +484 -0
- transport/utils.py +32 -0
transport/__init__.py
ADDED
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from .transport import Transport, ModelType, WeightType, PathType, SNRType, Sampler
|
2 |
+
|
3 |
+
|
4 |
+
def create_transport(
|
5 |
+
path_type="Linear",
|
6 |
+
prediction="velocity",
|
7 |
+
loss_weight=None,
|
8 |
+
train_eps=None,
|
9 |
+
sample_eps=None,
|
10 |
+
snr_type="uniform",
|
11 |
+
):
|
12 |
+
"""function for creating Transport object
|
13 |
+
**Note**: model prediction defaults to velocity
|
14 |
+
Args:
|
15 |
+
- path_type: type of path to use; default to linear
|
16 |
+
- learn_score: set model prediction to score
|
17 |
+
- learn_noise: set model prediction to noise
|
18 |
+
- velocity_weighted: weight loss by velocity weight
|
19 |
+
- likelihood_weighted: weight loss by likelihood weight
|
20 |
+
- train_eps: small epsilon for avoiding instability during training
|
21 |
+
- sample_eps: small epsilon for avoiding instability during sampling
|
22 |
+
"""
|
23 |
+
|
24 |
+
if prediction == "noise":
|
25 |
+
model_type = ModelType.NOISE
|
26 |
+
elif prediction == "score":
|
27 |
+
model_type = ModelType.SCORE
|
28 |
+
else:
|
29 |
+
model_type = ModelType.VELOCITY
|
30 |
+
|
31 |
+
if loss_weight == "velocity":
|
32 |
+
loss_type = WeightType.VELOCITY
|
33 |
+
elif loss_weight == "likelihood":
|
34 |
+
loss_type = WeightType.LIKELIHOOD
|
35 |
+
else:
|
36 |
+
loss_type = WeightType.NONE
|
37 |
+
|
38 |
+
if snr_type == "lognorm":
|
39 |
+
snr_type = SNRType.LOGNORM
|
40 |
+
elif snr_type == "uniform":
|
41 |
+
snr_type = SNRType.UNIFORM
|
42 |
+
else:
|
43 |
+
raise ValueError(f"Invalid snr type {snr_type}")
|
44 |
+
|
45 |
+
path_choice = {
|
46 |
+
"Linear": PathType.LINEAR,
|
47 |
+
"GVP": PathType.GVP,
|
48 |
+
"VP": PathType.VP,
|
49 |
+
}
|
50 |
+
|
51 |
+
path_type = path_choice[path_type]
|
52 |
+
|
53 |
+
if path_type in [PathType.VP]:
|
54 |
+
train_eps = 1e-5 if train_eps is None else train_eps
|
55 |
+
sample_eps = 1e-3 if train_eps is None else sample_eps
|
56 |
+
elif (
|
57 |
+
path_type in [PathType.GVP, PathType.LINEAR]
|
58 |
+
and model_type != ModelType.VELOCITY
|
59 |
+
):
|
60 |
+
train_eps = 1e-3 if train_eps is None else train_eps
|
61 |
+
sample_eps = 1e-3 if train_eps is None else sample_eps
|
62 |
+
else: # velocity & [GVP, LINEAR] is stable everywhere
|
63 |
+
train_eps = 0
|
64 |
+
sample_eps = 0
|
65 |
+
|
66 |
+
# create flow state
|
67 |
+
state = Transport(
|
68 |
+
model_type=model_type,
|
69 |
+
path_type=path_type,
|
70 |
+
loss_type=loss_type,
|
71 |
+
train_eps=train_eps,
|
72 |
+
sample_eps=sample_eps,
|
73 |
+
snr_type=snr_type,
|
74 |
+
)
|
75 |
+
|
76 |
+
return state
|
transport/integrators.py
ADDED
@@ -0,0 +1,127 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import torch as th
|
3 |
+
import torch.nn as nn
|
4 |
+
from torchdiffeq import odeint
|
5 |
+
from functools import partial
|
6 |
+
from tqdm import tqdm
|
7 |
+
|
8 |
+
|
9 |
+
class sde:
|
10 |
+
"""SDE solver class"""
|
11 |
+
|
12 |
+
def __init__(
|
13 |
+
self,
|
14 |
+
drift,
|
15 |
+
diffusion,
|
16 |
+
*,
|
17 |
+
t0,
|
18 |
+
t1,
|
19 |
+
num_steps,
|
20 |
+
sampler_type,
|
21 |
+
):
|
22 |
+
assert t0 < t1, "SDE sampler has to be in forward time"
|
23 |
+
|
24 |
+
self.num_timesteps = num_steps
|
25 |
+
self.t = th.linspace(t0, t1, num_steps)
|
26 |
+
self.dt = self.t[1] - self.t[0]
|
27 |
+
self.drift = drift
|
28 |
+
self.diffusion = diffusion
|
29 |
+
self.sampler_type = sampler_type
|
30 |
+
|
31 |
+
def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs):
|
32 |
+
w_cur = th.randn(x.size()).to(x)
|
33 |
+
t = th.ones(x.size(0)).to(x) * t
|
34 |
+
dw = w_cur * th.sqrt(self.dt)
|
35 |
+
drift = self.drift(x, t, model, **model_kwargs)
|
36 |
+
diffusion = self.diffusion(x, t)
|
37 |
+
mean_x = x + drift * self.dt
|
38 |
+
x = mean_x + th.sqrt(2 * diffusion) * dw
|
39 |
+
return x, mean_x
|
40 |
+
|
41 |
+
def __Heun_step(self, x, _, t, model, **model_kwargs):
|
42 |
+
w_cur = th.randn(x.size()).to(x)
|
43 |
+
dw = w_cur * th.sqrt(self.dt)
|
44 |
+
t_cur = th.ones(x.size(0)).to(x) * t
|
45 |
+
diffusion = self.diffusion(x, t_cur)
|
46 |
+
xhat = x + th.sqrt(2 * diffusion) * dw
|
47 |
+
K1 = self.drift(xhat, t_cur, model, **model_kwargs)
|
48 |
+
xp = xhat + self.dt * K1
|
49 |
+
K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs)
|
50 |
+
return (
|
51 |
+
xhat + 0.5 * self.dt * (K1 + K2),
|
52 |
+
xhat,
|
53 |
+
) # at last time point we do not perform the heun step
|
54 |
+
|
55 |
+
def __forward_fn(self):
|
56 |
+
"""TODO: generalize here by adding all private functions ending with steps to it"""
|
57 |
+
sampler_dict = {
|
58 |
+
"Euler": self.__Euler_Maruyama_step,
|
59 |
+
"Heun": self.__Heun_step,
|
60 |
+
}
|
61 |
+
|
62 |
+
try:
|
63 |
+
sampler = sampler_dict[self.sampler_type]
|
64 |
+
except:
|
65 |
+
raise NotImplementedError("Smapler type not implemented.")
|
66 |
+
|
67 |
+
return sampler
|
68 |
+
|
69 |
+
def sample(self, init, model, **model_kwargs):
|
70 |
+
"""forward loop of sde"""
|
71 |
+
x = init
|
72 |
+
mean_x = init
|
73 |
+
samples = []
|
74 |
+
sampler = self.__forward_fn()
|
75 |
+
for ti in self.t[:-1]:
|
76 |
+
with th.no_grad():
|
77 |
+
x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs)
|
78 |
+
samples.append(x)
|
79 |
+
|
80 |
+
return samples
|
81 |
+
|
82 |
+
|
83 |
+
class ode:
|
84 |
+
"""ODE solver class"""
|
85 |
+
|
86 |
+
def __init__(
|
87 |
+
self,
|
88 |
+
drift,
|
89 |
+
*,
|
90 |
+
t0,
|
91 |
+
t1,
|
92 |
+
sampler_type,
|
93 |
+
num_steps,
|
94 |
+
atol,
|
95 |
+
rtol,
|
96 |
+
time_shifting_factor=None,
|
97 |
+
):
|
98 |
+
assert t0 < t1, "ODE sampler has to be in forward time"
|
99 |
+
|
100 |
+
self.drift = drift
|
101 |
+
self.t = th.linspace(t0, t1, num_steps)
|
102 |
+
if time_shifting_factor:
|
103 |
+
self.t = self.t / (
|
104 |
+
self.t + time_shifting_factor - time_shifting_factor * self.t
|
105 |
+
)
|
106 |
+
self.atol = atol
|
107 |
+
self.rtol = rtol
|
108 |
+
self.sampler_type = sampler_type
|
109 |
+
|
110 |
+
def sample(self, x, model, **model_kwargs):
|
111 |
+
|
112 |
+
device = x[0].device if isinstance(x, tuple) else x.device
|
113 |
+
|
114 |
+
def _fn(t, x):
|
115 |
+
t = (
|
116 |
+
th.ones(x[0].size(0)).to(device) * t
|
117 |
+
if isinstance(x, tuple)
|
118 |
+
else th.ones(x.size(0)).to(device) * t
|
119 |
+
)
|
120 |
+
model_output = self.drift(x, t, model, **model_kwargs)
|
121 |
+
return model_output
|
122 |
+
|
123 |
+
t = self.t.to(device)
|
124 |
+
atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol]
|
125 |
+
rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol]
|
126 |
+
samples = odeint(_fn, x, t, method=self.sampler_type, atol=atol, rtol=rtol)
|
127 |
+
return samples
|
transport/path.py
ADDED
@@ -0,0 +1,206 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch as th
|
2 |
+
import numpy as np
|
3 |
+
from functools import partial
|
4 |
+
|
5 |
+
|
6 |
+
def expand_t_like_x(t, x):
|
7 |
+
"""Function to reshape time t to broadcastable dimension of x
|
8 |
+
Args:
|
9 |
+
t: [batch_dim,], time vector
|
10 |
+
x: [batch_dim,...], data point
|
11 |
+
"""
|
12 |
+
dims = [1] * len(x[0].size())
|
13 |
+
t = t.view(t.size(0), *dims)
|
14 |
+
return t
|
15 |
+
|
16 |
+
|
17 |
+
#################### Coupling Plans ####################
|
18 |
+
|
19 |
+
|
20 |
+
class ICPlan:
|
21 |
+
"""Linear Coupling Plan"""
|
22 |
+
|
23 |
+
def __init__(self, sigma=0.0):
|
24 |
+
self.sigma = sigma
|
25 |
+
|
26 |
+
def compute_alpha_t(self, t):
|
27 |
+
"""Compute the data coefficient along the path"""
|
28 |
+
return t, 1
|
29 |
+
|
30 |
+
def compute_sigma_t(self, t):
|
31 |
+
"""Compute the noise coefficient along the path"""
|
32 |
+
return 1 - t, -1
|
33 |
+
|
34 |
+
def compute_d_alpha_alpha_ratio_t(self, t):
|
35 |
+
"""Compute the ratio between d_alpha and alpha"""
|
36 |
+
return 1 / t
|
37 |
+
|
38 |
+
def compute_drift(self, x, t):
|
39 |
+
"""We always output sde according to score parametrization;"""
|
40 |
+
t = expand_t_like_x(t, x)
|
41 |
+
alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
|
42 |
+
sigma_t, d_sigma_t = self.compute_sigma_t(t)
|
43 |
+
drift = alpha_ratio * x
|
44 |
+
diffusion = alpha_ratio * (sigma_t**2) - sigma_t * d_sigma_t
|
45 |
+
|
46 |
+
return -drift, diffusion
|
47 |
+
|
48 |
+
def compute_diffusion(self, x, t, form="constant", norm=1.0):
|
49 |
+
"""Compute the diffusion term of the SDE
|
50 |
+
Args:
|
51 |
+
x: [batch_dim, ...], data point
|
52 |
+
t: [batch_dim,], time vector
|
53 |
+
form: str, form of the diffusion term
|
54 |
+
norm: float, norm of the diffusion term
|
55 |
+
"""
|
56 |
+
t = expand_t_like_x(t, x)
|
57 |
+
choices = {
|
58 |
+
"constant": norm,
|
59 |
+
"SBDM": norm * self.compute_drift(x, t)[1],
|
60 |
+
"sigma": norm * self.compute_sigma_t(t)[0],
|
61 |
+
"linear": norm * (1 - t),
|
62 |
+
"decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2,
|
63 |
+
"inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2,
|
64 |
+
}
|
65 |
+
|
66 |
+
try:
|
67 |
+
diffusion = choices[form]
|
68 |
+
except KeyError:
|
69 |
+
raise NotImplementedError(f"Diffusion form {form} not implemented")
|
70 |
+
|
71 |
+
return diffusion
|
72 |
+
|
73 |
+
def get_score_from_velocity(self, velocity, x, t):
|
74 |
+
"""Wrapper function: transfrom velocity prediction model to score
|
75 |
+
Args:
|
76 |
+
velocity: [batch_dim, ...] shaped tensor; velocity model output
|
77 |
+
x: [batch_dim, ...] shaped tensor; x_t data point
|
78 |
+
t: [batch_dim,] time tensor
|
79 |
+
"""
|
80 |
+
t = expand_t_like_x(t, x)
|
81 |
+
alpha_t, d_alpha_t = self.compute_alpha_t(t)
|
82 |
+
sigma_t, d_sigma_t = self.compute_sigma_t(t)
|
83 |
+
mean = x
|
84 |
+
reverse_alpha_ratio = alpha_t / d_alpha_t
|
85 |
+
var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
|
86 |
+
score = (reverse_alpha_ratio * velocity - mean) / var
|
87 |
+
return score
|
88 |
+
|
89 |
+
def get_noise_from_velocity(self, velocity, x, t):
|
90 |
+
"""Wrapper function: transfrom velocity prediction model to denoiser
|
91 |
+
Args:
|
92 |
+
velocity: [batch_dim, ...] shaped tensor; velocity model output
|
93 |
+
x: [batch_dim, ...] shaped tensor; x_t data point
|
94 |
+
t: [batch_dim,] time tensor
|
95 |
+
"""
|
96 |
+
t = expand_t_like_x(t, x)
|
97 |
+
alpha_t, d_alpha_t = self.compute_alpha_t(t)
|
98 |
+
sigma_t, d_sigma_t = self.compute_sigma_t(t)
|
99 |
+
mean = x
|
100 |
+
reverse_alpha_ratio = alpha_t / d_alpha_t
|
101 |
+
var = reverse_alpha_ratio * d_sigma_t - sigma_t
|
102 |
+
noise = (reverse_alpha_ratio * velocity - mean) / var
|
103 |
+
return noise
|
104 |
+
|
105 |
+
def get_velocity_from_score(self, score, x, t):
|
106 |
+
"""Wrapper function: transfrom score prediction model to velocity
|
107 |
+
Args:
|
108 |
+
score: [batch_dim, ...] shaped tensor; score model output
|
109 |
+
x: [batch_dim, ...] shaped tensor; x_t data point
|
110 |
+
t: [batch_dim,] time tensor
|
111 |
+
"""
|
112 |
+
t = expand_t_like_x(t, x)
|
113 |
+
drift, var = self.compute_drift(x, t)
|
114 |
+
velocity = var * score - drift
|
115 |
+
return velocity
|
116 |
+
|
117 |
+
def compute_mu_t(self, t, x0, x1):
|
118 |
+
"""Compute the mean of time-dependent density p_t"""
|
119 |
+
t = expand_t_like_x(t, x1)
|
120 |
+
alpha_t, _ = self.compute_alpha_t(t)
|
121 |
+
sigma_t, _ = self.compute_sigma_t(t)
|
122 |
+
if isinstance(x1, (list, tuple)):
|
123 |
+
return [alpha_t[i] * x1[i] + sigma_t[i] * x0[i] for i in range(len(x1))]
|
124 |
+
else:
|
125 |
+
return alpha_t * x1 + sigma_t * x0
|
126 |
+
|
127 |
+
def compute_xt(self, t, x0, x1):
|
128 |
+
"""Sample xt from time-dependent density p_t; rng is required"""
|
129 |
+
xt = self.compute_mu_t(t, x0, x1)
|
130 |
+
return xt
|
131 |
+
|
132 |
+
def compute_ut(self, t, x0, x1, xt):
|
133 |
+
"""Compute the vector field corresponding to p_t"""
|
134 |
+
t = expand_t_like_x(t, x1)
|
135 |
+
_, d_alpha_t = self.compute_alpha_t(t)
|
136 |
+
_, d_sigma_t = self.compute_sigma_t(t)
|
137 |
+
if isinstance(x1, (list, tuple)):
|
138 |
+
return [d_alpha_t * x1[i] + d_sigma_t * x0[i] for i in range(len(x1))]
|
139 |
+
else:
|
140 |
+
return d_alpha_t * x1 + d_sigma_t * x0
|
141 |
+
|
142 |
+
def plan(self, t, x0, x1):
|
143 |
+
xt = self.compute_xt(t, x0, x1)
|
144 |
+
ut = self.compute_ut(t, x0, x1, xt)
|
145 |
+
return t, xt, ut
|
146 |
+
|
147 |
+
|
148 |
+
class VPCPlan(ICPlan):
|
149 |
+
"""class for VP path flow matching"""
|
150 |
+
|
151 |
+
def __init__(self, sigma_min=0.1, sigma_max=20.0):
|
152 |
+
self.sigma_min = sigma_min
|
153 |
+
self.sigma_max = sigma_max
|
154 |
+
self.log_mean_coeff = (
|
155 |
+
lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min)
|
156 |
+
- 0.5 * (1 - t) * self.sigma_min
|
157 |
+
)
|
158 |
+
self.d_log_mean_coeff = (
|
159 |
+
lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min)
|
160 |
+
+ 0.5 * self.sigma_min
|
161 |
+
)
|
162 |
+
|
163 |
+
def compute_alpha_t(self, t):
|
164 |
+
"""Compute coefficient of x1"""
|
165 |
+
alpha_t = self.log_mean_coeff(t)
|
166 |
+
alpha_t = th.exp(alpha_t)
|
167 |
+
d_alpha_t = alpha_t * self.d_log_mean_coeff(t)
|
168 |
+
return alpha_t, d_alpha_t
|
169 |
+
|
170 |
+
def compute_sigma_t(self, t):
|
171 |
+
"""Compute coefficient of x0"""
|
172 |
+
p_sigma_t = 2 * self.log_mean_coeff(t)
|
173 |
+
sigma_t = th.sqrt(1 - th.exp(p_sigma_t))
|
174 |
+
d_sigma_t = th.exp(p_sigma_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * sigma_t)
|
175 |
+
return sigma_t, d_sigma_t
|
176 |
+
|
177 |
+
def compute_d_alpha_alpha_ratio_t(self, t):
|
178 |
+
"""Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
|
179 |
+
return self.d_log_mean_coeff(t)
|
180 |
+
|
181 |
+
def compute_drift(self, x, t):
|
182 |
+
"""Compute the drift term of the SDE"""
|
183 |
+
t = expand_t_like_x(t, x)
|
184 |
+
beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min)
|
185 |
+
return -0.5 * beta_t * x, beta_t / 2
|
186 |
+
|
187 |
+
|
188 |
+
class GVPCPlan(ICPlan):
|
189 |
+
def __init__(self, sigma=0.0):
|
190 |
+
super().__init__(sigma)
|
191 |
+
|
192 |
+
def compute_alpha_t(self, t):
|
193 |
+
"""Compute coefficient of x1"""
|
194 |
+
alpha_t = th.sin(t * np.pi / 2)
|
195 |
+
d_alpha_t = np.pi / 2 * th.cos(t * np.pi / 2)
|
196 |
+
return alpha_t, d_alpha_t
|
197 |
+
|
198 |
+
def compute_sigma_t(self, t):
|
199 |
+
"""Compute coefficient of x0"""
|
200 |
+
sigma_t = th.cos(t * np.pi / 2)
|
201 |
+
d_sigma_t = -np.pi / 2 * th.sin(t * np.pi / 2)
|
202 |
+
return sigma_t, d_sigma_t
|
203 |
+
|
204 |
+
def compute_d_alpha_alpha_ratio_t(self, t):
|
205 |
+
"""Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
|
206 |
+
return np.pi / (2 * th.tan(t * np.pi / 2))
|
transport/transport.py
ADDED
@@ -0,0 +1,484 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch as th
|
2 |
+
import numpy as np
|
3 |
+
import logging
|
4 |
+
|
5 |
+
import enum
|
6 |
+
|
7 |
+
from . import path
|
8 |
+
from .utils import EasyDict, log_state, mean_flat
|
9 |
+
from .integrators import ode, sde
|
10 |
+
|
11 |
+
|
12 |
+
class ModelType(enum.Enum):
|
13 |
+
"""
|
14 |
+
Which type of output the model predicts.
|
15 |
+
"""
|
16 |
+
|
17 |
+
NOISE = enum.auto() # the model predicts epsilon
|
18 |
+
SCORE = enum.auto() # the model predicts \nabla \log p(x)
|
19 |
+
VELOCITY = enum.auto() # the model predicts v(x)
|
20 |
+
|
21 |
+
|
22 |
+
class PathType(enum.Enum):
|
23 |
+
"""
|
24 |
+
Which type of path to use.
|
25 |
+
"""
|
26 |
+
|
27 |
+
LINEAR = enum.auto()
|
28 |
+
GVP = enum.auto()
|
29 |
+
VP = enum.auto()
|
30 |
+
|
31 |
+
|
32 |
+
class WeightType(enum.Enum):
|
33 |
+
"""
|
34 |
+
Which type of weighting to use.
|
35 |
+
"""
|
36 |
+
|
37 |
+
NONE = enum.auto()
|
38 |
+
VELOCITY = enum.auto()
|
39 |
+
LIKELIHOOD = enum.auto()
|
40 |
+
|
41 |
+
|
42 |
+
class SNRType(enum.Enum):
|
43 |
+
UNIFORM = enum.auto()
|
44 |
+
LOGNORM = enum.auto()
|
45 |
+
|
46 |
+
|
47 |
+
class Transport:
|
48 |
+
|
49 |
+
def __init__(
|
50 |
+
self, *, model_type, path_type, loss_type, train_eps, sample_eps, snr_type
|
51 |
+
):
|
52 |
+
path_options = {
|
53 |
+
PathType.LINEAR: path.ICPlan,
|
54 |
+
PathType.GVP: path.GVPCPlan,
|
55 |
+
PathType.VP: path.VPCPlan,
|
56 |
+
}
|
57 |
+
|
58 |
+
self.loss_type = loss_type
|
59 |
+
self.model_type = model_type
|
60 |
+
self.path_sampler = path_options[path_type]()
|
61 |
+
self.train_eps = train_eps
|
62 |
+
self.sample_eps = sample_eps
|
63 |
+
|
64 |
+
self.snr_type = snr_type
|
65 |
+
|
66 |
+
def prior_logp(self, z):
|
67 |
+
"""
|
68 |
+
Standard multivariate normal prior
|
69 |
+
Assume z is batched
|
70 |
+
"""
|
71 |
+
shape = th.tensor(z.size())
|
72 |
+
N = th.prod(shape[1:])
|
73 |
+
_fn = lambda x: -N / 2.0 * np.log(2 * np.pi) - th.sum(x**2) / 2.0
|
74 |
+
return th.vmap(_fn)(z)
|
75 |
+
|
76 |
+
def check_interval(
|
77 |
+
self,
|
78 |
+
train_eps,
|
79 |
+
sample_eps,
|
80 |
+
*,
|
81 |
+
diffusion_form="SBDM",
|
82 |
+
sde=False,
|
83 |
+
reverse=False,
|
84 |
+
eval=False,
|
85 |
+
last_step_size=0.0,
|
86 |
+
):
|
87 |
+
t0 = 0
|
88 |
+
t1 = 1
|
89 |
+
eps = train_eps if not eval else sample_eps
|
90 |
+
if type(self.path_sampler) in [path.VPCPlan]:
|
91 |
+
|
92 |
+
t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
|
93 |
+
|
94 |
+
elif (type(self.path_sampler) in [path.ICPlan, path.GVPCPlan]) and (
|
95 |
+
self.model_type != ModelType.VELOCITY or sde
|
96 |
+
): # avoid numerical issue by taking a first semi-implicit step
|
97 |
+
|
98 |
+
t0 = (
|
99 |
+
eps
|
100 |
+
if (diffusion_form == "SBDM" and sde)
|
101 |
+
or self.model_type != ModelType.VELOCITY
|
102 |
+
else 0
|
103 |
+
)
|
104 |
+
t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
|
105 |
+
|
106 |
+
if reverse:
|
107 |
+
t0, t1 = 1 - t0, 1 - t1
|
108 |
+
|
109 |
+
return t0, t1
|
110 |
+
|
111 |
+
def sample(self, x1):
|
112 |
+
"""Sampling x0 & t based on shape of x1 (if needed)
|
113 |
+
Args:
|
114 |
+
x1 - data point; [batch, *dim]
|
115 |
+
"""
|
116 |
+
if isinstance(x1, (list, tuple)):
|
117 |
+
x0 = [th.randn_like(img_start) for img_start in x1]
|
118 |
+
else:
|
119 |
+
x0 = th.randn_like(x1)
|
120 |
+
t0, t1 = self.check_interval(self.train_eps, self.sample_eps)
|
121 |
+
|
122 |
+
if self.snr_type == SNRType.UNIFORM:
|
123 |
+
t = th.rand((len(x1),)) * (t1 - t0) + t0
|
124 |
+
elif self.snr_type == SNRType.LOGNORM:
|
125 |
+
u = th.normal(mean=0.0, std=1.0, size=(len(x1),))
|
126 |
+
t = 1 / (1 + th.exp(-u)) * (t1 - t0) + t0
|
127 |
+
else:
|
128 |
+
raise ValueError(f"Unknown snr type: {self.snr_type}")
|
129 |
+
t = t.to(x1[0])
|
130 |
+
return t, x0, x1
|
131 |
+
|
132 |
+
def training_losses(self, model, x1, model_kwargs=None):
|
133 |
+
"""Loss for training the score model
|
134 |
+
Args:
|
135 |
+
- model: backbone model; could be score, noise, or velocity
|
136 |
+
- x1: datapoint
|
137 |
+
- model_kwargs: additional arguments for the model
|
138 |
+
"""
|
139 |
+
if model_kwargs == None:
|
140 |
+
model_kwargs = {}
|
141 |
+
t, x0, x1 = self.sample(x1)
|
142 |
+
t, xt, ut = self.path_sampler.plan(t, x0, x1)
|
143 |
+
model_output = model(xt, t, **model_kwargs)
|
144 |
+
B = len(x0)
|
145 |
+
|
146 |
+
terms = {}
|
147 |
+
# terms['pred'] = model_output
|
148 |
+
if self.model_type == ModelType.VELOCITY:
|
149 |
+
if isinstance(x1, (list, tuple)):
|
150 |
+
assert len(model_output) == len(ut) == len(x1)
|
151 |
+
for i in range(B):
|
152 |
+
assert (
|
153 |
+
model_output[i].shape == ut[i].shape == x1[i].shape
|
154 |
+
), f"{model_output[i].shape} {ut[i].shape} {x1[i].shape}"
|
155 |
+
terms["task_loss"] = th.stack(
|
156 |
+
[((ut[i] - model_output[i]) ** 2).mean() for i in range(B)],
|
157 |
+
dim=0,
|
158 |
+
)
|
159 |
+
else:
|
160 |
+
terms["task_loss"] = mean_flat(((model_output - ut) ** 2))
|
161 |
+
else:
|
162 |
+
raise NotImplementedError
|
163 |
+
# _, drift_var = self.path_sampler.compute_drift(xt, t)
|
164 |
+
# sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, xt))
|
165 |
+
# if self.loss_type in [WeightType.VELOCITY]:
|
166 |
+
# weight = (drift_var / sigma_t) ** 2
|
167 |
+
# elif self.loss_type in [WeightType.LIKELIHOOD]:
|
168 |
+
# weight = drift_var / (sigma_t ** 2)
|
169 |
+
# elif self.loss_type in [WeightType.NONE]:
|
170 |
+
# weight = 1
|
171 |
+
# else:
|
172 |
+
# raise NotImplementedError()
|
173 |
+
#
|
174 |
+
# if self.model_type == ModelType.NOISE:
|
175 |
+
# terms['task_loss'] = mean_flat(weight * ((model_output - x0) ** 2))
|
176 |
+
# else:
|
177 |
+
# terms['task_loss'] = mean_flat(weight * ((model_output * sigma_t + x0) ** 2))
|
178 |
+
|
179 |
+
terms["loss"] = terms["task_loss"]
|
180 |
+
terms["task_loss"] = terms["task_loss"].clone().detach()
|
181 |
+
return terms
|
182 |
+
|
183 |
+
def get_drift(self):
|
184 |
+
"""member function for obtaining the drift of the probability flow ODE"""
|
185 |
+
|
186 |
+
def score_ode(x, t, model, **model_kwargs):
|
187 |
+
drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
|
188 |
+
model_output = model(x, t, **model_kwargs)
|
189 |
+
return -drift_mean + drift_var * model_output # by change of variable
|
190 |
+
|
191 |
+
def noise_ode(x, t, model, **model_kwargs):
|
192 |
+
drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
|
193 |
+
sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))
|
194 |
+
model_output = model(x, t, **model_kwargs)
|
195 |
+
score = model_output / -sigma_t
|
196 |
+
return -drift_mean + drift_var * score
|
197 |
+
|
198 |
+
def velocity_ode(x, t, model, **model_kwargs):
|
199 |
+
model_output = model(x, t, **model_kwargs)
|
200 |
+
return model_output
|
201 |
+
|
202 |
+
if self.model_type == ModelType.NOISE:
|
203 |
+
drift_fn = noise_ode
|
204 |
+
elif self.model_type == ModelType.SCORE:
|
205 |
+
drift_fn = score_ode
|
206 |
+
else:
|
207 |
+
drift_fn = velocity_ode
|
208 |
+
|
209 |
+
def body_fn(x, t, model, **model_kwargs):
|
210 |
+
model_output = drift_fn(x, t, model, **model_kwargs)
|
211 |
+
assert (
|
212 |
+
model_output.shape == x.shape
|
213 |
+
), "Output shape from ODE solver must match input shape"
|
214 |
+
return model_output
|
215 |
+
|
216 |
+
return body_fn
|
217 |
+
|
218 |
+
def get_score(
|
219 |
+
self,
|
220 |
+
):
|
221 |
+
"""member function for obtaining score of
|
222 |
+
x_t = alpha_t * x + sigma_t * eps"""
|
223 |
+
if self.model_type == ModelType.NOISE:
|
224 |
+
score_fn = (
|
225 |
+
lambda x, t, model, **kwargs: model(x, t, **kwargs)
|
226 |
+
/ -self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))[0]
|
227 |
+
)
|
228 |
+
elif self.model_type == ModelType.SCORE:
|
229 |
+
score_fn = lambda x, t, model, **kwagrs: model(x, t, **kwagrs)
|
230 |
+
elif self.model_type == ModelType.VELOCITY:
|
231 |
+
score_fn = (
|
232 |
+
lambda x, t, model, **kwargs: self.path_sampler.get_score_from_velocity(
|
233 |
+
model(x, t, **kwargs), x, t
|
234 |
+
)
|
235 |
+
)
|
236 |
+
else:
|
237 |
+
raise NotImplementedError()
|
238 |
+
|
239 |
+
return score_fn
|
240 |
+
|
241 |
+
|
242 |
+
class Sampler:
|
243 |
+
"""Sampler class for the transport model"""
|
244 |
+
|
245 |
+
def __init__(
|
246 |
+
self,
|
247 |
+
transport,
|
248 |
+
):
|
249 |
+
"""Constructor for a general sampler; supporting different sampling methods
|
250 |
+
Args:
|
251 |
+
- transport: an tranport object specify model prediction & interpolant type
|
252 |
+
"""
|
253 |
+
|
254 |
+
self.transport = transport
|
255 |
+
self.drift = self.transport.get_drift()
|
256 |
+
self.score = self.transport.get_score()
|
257 |
+
|
258 |
+
def __get_sde_diffusion_and_drift(
|
259 |
+
self,
|
260 |
+
*,
|
261 |
+
diffusion_form="SBDM",
|
262 |
+
diffusion_norm=1.0,
|
263 |
+
):
|
264 |
+
|
265 |
+
def diffusion_fn(x, t):
|
266 |
+
diffusion = self.transport.path_sampler.compute_diffusion(
|
267 |
+
x, t, form=diffusion_form, norm=diffusion_norm
|
268 |
+
)
|
269 |
+
return diffusion
|
270 |
+
|
271 |
+
sde_drift = lambda x, t, model, **kwargs: self.drift(
|
272 |
+
x, t, model, **kwargs
|
273 |
+
) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs)
|
274 |
+
|
275 |
+
sde_diffusion = diffusion_fn
|
276 |
+
|
277 |
+
return sde_drift, sde_diffusion
|
278 |
+
|
279 |
+
def __get_last_step(
|
280 |
+
self,
|
281 |
+
sde_drift,
|
282 |
+
*,
|
283 |
+
last_step,
|
284 |
+
last_step_size,
|
285 |
+
):
|
286 |
+
"""Get the last step function of the SDE solver"""
|
287 |
+
|
288 |
+
if last_step is None:
|
289 |
+
last_step_fn = lambda x, t, model, **model_kwargs: x
|
290 |
+
elif last_step == "Mean":
|
291 |
+
last_step_fn = (
|
292 |
+
lambda x, t, model, **model_kwargs: x
|
293 |
+
+ sde_drift(x, t, model, **model_kwargs) * last_step_size
|
294 |
+
)
|
295 |
+
elif last_step == "Tweedie":
|
296 |
+
alpha = (
|
297 |
+
self.transport.path_sampler.compute_alpha_t
|
298 |
+
) # simple aliasing; the original name was too long
|
299 |
+
sigma = self.transport.path_sampler.compute_sigma_t
|
300 |
+
last_step_fn = lambda x, t, model, **model_kwargs: x / alpha(t)[0][0] + (
|
301 |
+
sigma(t)[0][0] ** 2
|
302 |
+
) / alpha(t)[0][0] * self.score(x, t, model, **model_kwargs)
|
303 |
+
elif last_step == "Euler":
|
304 |
+
last_step_fn = (
|
305 |
+
lambda x, t, model, **model_kwargs: x
|
306 |
+
+ self.drift(x, t, model, **model_kwargs) * last_step_size
|
307 |
+
)
|
308 |
+
else:
|
309 |
+
raise NotImplementedError()
|
310 |
+
|
311 |
+
return last_step_fn
|
312 |
+
|
313 |
+
def sample_sde(
|
314 |
+
self,
|
315 |
+
*,
|
316 |
+
sampling_method="Euler",
|
317 |
+
diffusion_form="SBDM",
|
318 |
+
diffusion_norm=1.0,
|
319 |
+
last_step="Mean",
|
320 |
+
last_step_size=0.04,
|
321 |
+
num_steps=250,
|
322 |
+
):
|
323 |
+
"""returns a sampling function with given SDE settings
|
324 |
+
Args:
|
325 |
+
- sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama
|
326 |
+
- diffusion_form: function form of diffusion coefficient; default to be matching SBDM
|
327 |
+
- diffusion_norm: function magnitude of diffusion coefficient; default to 1
|
328 |
+
- last_step: type of the last step; default to identity
|
329 |
+
- last_step_size: size of the last step; default to match the stride of 250 steps over [0,1]
|
330 |
+
- num_steps: total integration step of SDE
|
331 |
+
"""
|
332 |
+
|
333 |
+
if last_step is None:
|
334 |
+
last_step_size = 0.0
|
335 |
+
|
336 |
+
sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift(
|
337 |
+
diffusion_form=diffusion_form,
|
338 |
+
diffusion_norm=diffusion_norm,
|
339 |
+
)
|
340 |
+
|
341 |
+
t0, t1 = self.transport.check_interval(
|
342 |
+
self.transport.train_eps,
|
343 |
+
self.transport.sample_eps,
|
344 |
+
diffusion_form=diffusion_form,
|
345 |
+
sde=True,
|
346 |
+
eval=True,
|
347 |
+
reverse=False,
|
348 |
+
last_step_size=last_step_size,
|
349 |
+
)
|
350 |
+
|
351 |
+
_sde = sde(
|
352 |
+
sde_drift,
|
353 |
+
sde_diffusion,
|
354 |
+
t0=t0,
|
355 |
+
t1=t1,
|
356 |
+
num_steps=num_steps,
|
357 |
+
sampler_type=sampling_method,
|
358 |
+
)
|
359 |
+
|
360 |
+
last_step_fn = self.__get_last_step(
|
361 |
+
sde_drift, last_step=last_step, last_step_size=last_step_size
|
362 |
+
)
|
363 |
+
|
364 |
+
def _sample(init, model, **model_kwargs):
|
365 |
+
xs = _sde.sample(init, model, **model_kwargs)
|
366 |
+
ts = th.ones(init.size(0), device=init.device) * t1
|
367 |
+
x = last_step_fn(xs[-1], ts, model, **model_kwargs)
|
368 |
+
xs.append(x)
|
369 |
+
|
370 |
+
assert len(xs) == num_steps, "Samples does not match the number of steps"
|
371 |
+
|
372 |
+
return xs
|
373 |
+
|
374 |
+
return _sample
|
375 |
+
|
376 |
+
def sample_ode(
|
377 |
+
self,
|
378 |
+
*,
|
379 |
+
sampling_method="dopri5",
|
380 |
+
num_steps=50,
|
381 |
+
atol=1e-6,
|
382 |
+
rtol=1e-3,
|
383 |
+
reverse=False,
|
384 |
+
time_shifting_factor=None,
|
385 |
+
):
|
386 |
+
"""returns a sampling function with given ODE settings
|
387 |
+
Args:
|
388 |
+
- sampling_method: type of sampler used in solving the ODE; default to be Dopri5
|
389 |
+
- num_steps:
|
390 |
+
- fixed solver (Euler, Heun): the actual number of integration steps performed
|
391 |
+
- adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
|
392 |
+
- atol: absolute error tolerance for the solver
|
393 |
+
- rtol: relative error tolerance for the solver
|
394 |
+
- reverse: whether solving the ODE in reverse (data to noise); default to False
|
395 |
+
"""
|
396 |
+
if reverse:
|
397 |
+
drift = lambda x, t, model, **kwargs: self.drift(
|
398 |
+
x, th.ones_like(t) * (1 - t), model, **kwargs
|
399 |
+
)
|
400 |
+
else:
|
401 |
+
drift = self.drift
|
402 |
+
|
403 |
+
t0, t1 = self.transport.check_interval(
|
404 |
+
self.transport.train_eps,
|
405 |
+
self.transport.sample_eps,
|
406 |
+
sde=False,
|
407 |
+
eval=True,
|
408 |
+
reverse=reverse,
|
409 |
+
last_step_size=0.0,
|
410 |
+
)
|
411 |
+
|
412 |
+
_ode = ode(
|
413 |
+
drift=drift,
|
414 |
+
t0=t0,
|
415 |
+
t1=t1,
|
416 |
+
sampler_type=sampling_method,
|
417 |
+
num_steps=num_steps,
|
418 |
+
atol=atol,
|
419 |
+
rtol=rtol,
|
420 |
+
time_shifting_factor=time_shifting_factor,
|
421 |
+
)
|
422 |
+
|
423 |
+
return _ode.sample
|
424 |
+
|
425 |
+
def sample_ode_likelihood(
|
426 |
+
self,
|
427 |
+
*,
|
428 |
+
sampling_method="dopri5",
|
429 |
+
num_steps=50,
|
430 |
+
atol=1e-6,
|
431 |
+
rtol=1e-3,
|
432 |
+
):
|
433 |
+
"""returns a sampling function for calculating likelihood with given ODE settings
|
434 |
+
Args:
|
435 |
+
- sampling_method: type of sampler used in solving the ODE; default to be Dopri5
|
436 |
+
- num_steps:
|
437 |
+
- fixed solver (Euler, Heun): the actual number of integration steps performed
|
438 |
+
- adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
|
439 |
+
- atol: absolute error tolerance for the solver
|
440 |
+
- rtol: relative error tolerance for the solver
|
441 |
+
"""
|
442 |
+
|
443 |
+
def _likelihood_drift(x, t, model, **model_kwargs):
|
444 |
+
x, _ = x
|
445 |
+
eps = th.randint(2, x.size(), dtype=th.float, device=x.device) * 2 - 1
|
446 |
+
t = th.ones_like(t) * (1 - t)
|
447 |
+
with th.enable_grad():
|
448 |
+
x.requires_grad = True
|
449 |
+
grad = th.autograd.grad(
|
450 |
+
th.sum(self.drift(x, t, model, **model_kwargs) * eps), x
|
451 |
+
)[0]
|
452 |
+
logp_grad = th.sum(grad * eps, dim=tuple(range(1, len(x.size()))))
|
453 |
+
drift = self.drift(x, t, model, **model_kwargs)
|
454 |
+
return (-drift, logp_grad)
|
455 |
+
|
456 |
+
t0, t1 = self.transport.check_interval(
|
457 |
+
self.transport.train_eps,
|
458 |
+
self.transport.sample_eps,
|
459 |
+
sde=False,
|
460 |
+
eval=True,
|
461 |
+
reverse=False,
|
462 |
+
last_step_size=0.0,
|
463 |
+
)
|
464 |
+
|
465 |
+
_ode = ode(
|
466 |
+
drift=_likelihood_drift,
|
467 |
+
t0=t0,
|
468 |
+
t1=t1,
|
469 |
+
sampler_type=sampling_method,
|
470 |
+
num_steps=num_steps,
|
471 |
+
atol=atol,
|
472 |
+
rtol=rtol,
|
473 |
+
)
|
474 |
+
|
475 |
+
def _sample_fn(x, model, **model_kwargs):
|
476 |
+
init_logp = th.zeros(x.size(0)).to(x)
|
477 |
+
input = (x, init_logp)
|
478 |
+
drift, delta_logp = _ode.sample(input, model, **model_kwargs)
|
479 |
+
drift, delta_logp = drift[-1], delta_logp[-1]
|
480 |
+
prior_logp = self.transport.prior_logp(drift)
|
481 |
+
logp = prior_logp - delta_logp
|
482 |
+
return logp, drift
|
483 |
+
|
484 |
+
return _sample_fn
|
transport/utils.py
ADDED
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch as th
|
2 |
+
|
3 |
+
|
4 |
+
class EasyDict:
|
5 |
+
|
6 |
+
def __init__(self, sub_dict):
|
7 |
+
for k, v in sub_dict.items():
|
8 |
+
setattr(self, k, v)
|
9 |
+
|
10 |
+
def __getitem__(self, key):
|
11 |
+
return getattr(self, key)
|
12 |
+
|
13 |
+
|
14 |
+
def mean_flat(x):
|
15 |
+
"""
|
16 |
+
Take the mean over all non-batch dimensions.
|
17 |
+
"""
|
18 |
+
return th.mean(x, dim=list(range(1, len(x.size()))))
|
19 |
+
|
20 |
+
|
21 |
+
def log_state(state):
|
22 |
+
result = []
|
23 |
+
|
24 |
+
sorted_state = dict(sorted(state.items()))
|
25 |
+
for key, value in sorted_state.items():
|
26 |
+
# Check if the value is an instance of a class
|
27 |
+
if "<object" in str(value) or "object at" in str(value):
|
28 |
+
result.append(f"{key}: [{value.__class__.__name__}]")
|
29 |
+
else:
|
30 |
+
result.append(f"{key}: {value}")
|
31 |
+
|
32 |
+
return "\n".join(result)
|