update_name
This commit is contained in:
0
graph_dit/diffusion/__init__.py
Normal file
0
graph_dit/diffusion/__init__.py
Normal file
224
graph_dit/diffusion/diffusion_utils.py
Normal file
224
graph_dit/diffusion/diffusion_utils.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
import numpy as np
|
||||
from utils import PlaceHolder
|
||||
|
||||
|
||||
def sum_except_batch(x):
|
||||
return x.reshape(x.size(0), -1).sum(dim=-1)
|
||||
|
||||
def assert_correctly_masked(variable, node_mask):
|
||||
assert (variable * (1 - node_mask.long())).abs().max().item() < 1e-4, \
|
||||
'Variables not masked properly.'
|
||||
|
||||
def cosine_beta_schedule_discrete(timesteps, s=0.008):
|
||||
""" Cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ. """
|
||||
steps = timesteps + 2
|
||||
x = np.linspace(0, steps, steps)
|
||||
|
||||
alphas_cumprod = np.cos(0.5 * np.pi * ((x / steps) + s) / (1 + s)) ** 2
|
||||
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
||||
alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])
|
||||
betas = 1 - alphas
|
||||
return betas.squeeze()
|
||||
|
||||
def custom_beta_schedule_discrete(timesteps, average_num_nodes=30, s=0.008):
|
||||
""" Cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ. """
|
||||
steps = timesteps + 2
|
||||
x = np.linspace(0, steps, steps)
|
||||
|
||||
alphas_cumprod = np.cos(0.5 * np.pi * ((x / steps) + s) / (1 + s)) ** 2
|
||||
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
||||
alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])
|
||||
betas = 1 - alphas
|
||||
|
||||
assert timesteps >= 100
|
||||
|
||||
p = 4 / 5 # 1 - 1 / num_edge_classes
|
||||
num_edges = average_num_nodes * (average_num_nodes - 1) / 2
|
||||
|
||||
# First 100 steps: only a few updates per graph
|
||||
updates_per_graph = 1.2
|
||||
beta_first = updates_per_graph / (p * num_edges)
|
||||
|
||||
betas[betas < beta_first] = beta_first
|
||||
return np.array(betas)
|
||||
|
||||
|
||||
def check_mask_correct(variables, node_mask):
|
||||
for i, variable in enumerate(variables):
|
||||
if len(variable) > 0:
|
||||
assert_correctly_masked(variable, node_mask)
|
||||
|
||||
|
||||
def check_tensor_same_size(*args):
|
||||
for i, arg in enumerate(args):
|
||||
if i == 0:
|
||||
continue
|
||||
assert args[0].size() == arg.size()
|
||||
|
||||
|
||||
|
||||
def reverse_tensor(x):
|
||||
return x[torch.arange(x.size(0) - 1, -1, -1)]
|
||||
|
||||
|
||||
def sample_discrete_features(probX, probE, node_mask, step=None, add_nose=True):
|
||||
''' Sample features from multinomial distribution with given probabilities (probX, probE, proby)
|
||||
:param probX: bs, n, dx_out node features
|
||||
:param probE: bs, n, n, de_out edge features
|
||||
:param proby: bs, dy_out global features.
|
||||
'''
|
||||
bs, n, _ = probX.shape
|
||||
|
||||
# Noise X
|
||||
# The masked rows should define probability distributions as well
|
||||
probX[~node_mask] = 1 / probX.shape[-1]
|
||||
|
||||
# Flatten the probability tensor to sample with multinomial
|
||||
probX = probX.reshape(bs * n, -1) # (bs * n, dx_out)
|
||||
|
||||
# Sample X
|
||||
probX = probX + 1e-12
|
||||
probX = probX / probX.sum(dim=-1, keepdim=True)
|
||||
X_t = probX.multinomial(1) # (bs * n, 1)
|
||||
X_t = X_t.reshape(bs, n) # (bs, n)
|
||||
|
||||
# Noise E
|
||||
# The masked rows should define probability distributions as well
|
||||
inverse_edge_mask = ~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2))
|
||||
diag_mask = torch.eye(n).unsqueeze(0).expand(bs, -1, -1)
|
||||
|
||||
probE[inverse_edge_mask] = 1 / probE.shape[-1]
|
||||
probE[diag_mask.bool()] = 1 / probE.shape[-1]
|
||||
probE = probE.reshape(bs * n * n, -1) # (bs * n * n, de_out)
|
||||
probE = probE + 1e-12
|
||||
probE = probE / probE.sum(dim=-1, keepdim=True)
|
||||
|
||||
# Sample E
|
||||
E_t = probE.multinomial(1).reshape(bs, n, n) # (bs, n, n)
|
||||
E_t = torch.triu(E_t, diagonal=1)
|
||||
E_t = (E_t + torch.transpose(E_t, 1, 2))
|
||||
|
||||
return PlaceHolder(X=X_t, E=E_t, y=torch.zeros(bs, 0).type_as(X_t))
|
||||
|
||||
|
||||
def compute_batched_over0_posterior_distribution(X_t, Qt, Qsb, Qtb):
|
||||
""" M: X or E
|
||||
Compute xt @ Qt.T * x0 @ Qsb / x0 @ Qtb @ xt.T for each possible value of x0
|
||||
X_t: bs, n, dt or bs, n, n, dt
|
||||
Qt: bs, d_t-1, dt
|
||||
Qsb: bs, d0, d_t-1
|
||||
Qtb: bs, d0, dt.
|
||||
"""
|
||||
X_t = X_t.float()
|
||||
Qt_T = Qt.transpose(-1, -2).float() # bs, N, dt
|
||||
assert Qt.dim() == 3
|
||||
left_term = X_t @ Qt_T
|
||||
left_term = left_term.unsqueeze(dim=2) # bs, N, 1, d_t-1
|
||||
right_term = Qsb.unsqueeze(1)
|
||||
numerator = left_term * right_term # bs, N, d0, d_t-1
|
||||
|
||||
denominator = Qtb @ X_t.transpose(-1, -2) # bs, d0, N
|
||||
denominator = denominator.transpose(-1, -2) # bs, N, d0
|
||||
denominator = denominator.unsqueeze(-1) # bs, N, d0, 1
|
||||
|
||||
denominator[denominator == 0] = 1.
|
||||
return numerator / denominator
|
||||
|
||||
|
||||
def mask_distributions(true_X, true_E, pred_X, pred_E, node_mask):
|
||||
# Add a small value everywhere to avoid nans
|
||||
pred_X = pred_X.clamp_min(1e-18)
|
||||
pred_X = pred_X / torch.sum(pred_X, dim=-1, keepdim=True)
|
||||
|
||||
pred_E = pred_E.clamp_min(1e-18)
|
||||
pred_E = pred_E / torch.sum(pred_E, dim=-1, keepdim=True)
|
||||
|
||||
# Set masked rows to arbitrary distributions, so it doesn't contribute to loss
|
||||
row_X = torch.ones(true_X.size(-1), dtype=true_X.dtype, device=true_X.device)
|
||||
row_E = torch.zeros(true_E.size(-1), dtype=true_E.dtype, device=true_E.device).clamp_min(1e-18)
|
||||
row_E[0] = 1.
|
||||
|
||||
diag_mask = ~torch.eye(node_mask.size(1), device=node_mask.device, dtype=torch.bool).unsqueeze(0)
|
||||
true_X[~node_mask] = row_X
|
||||
true_E[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2) * diag_mask), :] = row_E
|
||||
pred_X[~node_mask] = row_X
|
||||
pred_E[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2) * diag_mask), :] = row_E
|
||||
|
||||
return true_X, true_E, pred_X, pred_E
|
||||
|
||||
def posterior_distributions(X, X_t, Qt, Qsb, Qtb, X_dim):
|
||||
bs, n, d = X.shape
|
||||
X = X.float()
|
||||
Qt_X_T = torch.transpose(Qt.X, -2, -1).float() # (bs, d, d)
|
||||
left_term = X_t @ Qt_X_T # (bs, N, d)
|
||||
right_term = X @ Qsb.X # (bs, N, d)
|
||||
|
||||
numerator = left_term * right_term # (bs, N, d)
|
||||
denominator = X @ Qtb.X # (bs, N, d) @ (bs, d, d) = (bs, N, d)
|
||||
denominator = denominator * X_t
|
||||
|
||||
num_X = numerator[:, :, :X_dim]
|
||||
num_E = numerator[:, :, X_dim:].reshape(bs, n*n, -1)
|
||||
|
||||
deno_X = denominator[:, :, :X_dim]
|
||||
deno_E = denominator[:, :, X_dim:].reshape(bs, n*n, -1)
|
||||
|
||||
# denominator = (denominator * X_t).sum(dim=-1) # (bs, N, d) * (bs, N, d) + sum = (bs, N)
|
||||
denominator = denominator.unsqueeze(-1) # (bs, N, 1)
|
||||
|
||||
deno_X = deno_X.sum(dim=-1).unsqueeze(-1)
|
||||
deno_E = deno_E.sum(dim=-1).unsqueeze(-1)
|
||||
|
||||
deno_X[deno_X == 0.] = 1
|
||||
deno_E[deno_E == 0.] = 1
|
||||
prob_X = num_X / deno_X
|
||||
prob_E = num_E / deno_E
|
||||
|
||||
prob_E = prob_E / prob_E.sum(dim=-1, keepdim=True)
|
||||
prob_X = prob_X / prob_X.sum(dim=-1, keepdim=True)
|
||||
return PlaceHolder(X=prob_X, E=prob_E, y=None)
|
||||
|
||||
|
||||
def sample_discrete_feature_noise(limit_dist, node_mask):
|
||||
""" Sample from the limit distribution of the diffusion process"""
|
||||
bs, n_max = node_mask.shape
|
||||
x_limit = limit_dist.X[None, None, :].expand(bs, n_max, -1)
|
||||
U_X = x_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max)
|
||||
U_X = F.one_hot(U_X.long(), num_classes=x_limit.shape[-1]).float()
|
||||
|
||||
e_limit = limit_dist.E[None, None, None, :].expand(bs, n_max, n_max, -1)
|
||||
U_E = e_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max, n_max)
|
||||
U_E = F.one_hot(U_E.long(), num_classes=e_limit.shape[-1]).float()
|
||||
|
||||
U_X = U_X.to(node_mask.device)
|
||||
U_E = U_E.to(node_mask.device)
|
||||
|
||||
# Get upper triangular part of edge noise, without main diagonal
|
||||
upper_triangular_mask = torch.zeros_like(U_E)
|
||||
indices = torch.triu_indices(row=U_E.size(1), col=U_E.size(2), offset=1)
|
||||
upper_triangular_mask[:, indices[0], indices[1], :] = 1
|
||||
|
||||
U_E = U_E * upper_triangular_mask
|
||||
U_E = (U_E + torch.transpose(U_E, 1, 2))
|
||||
|
||||
assert (U_E == torch.transpose(U_E, 1, 2)).all()
|
||||
return PlaceHolder(X=U_X, E=U_E, y=None).mask(node_mask)
|
||||
|
||||
def index_QE(X, q_e, n_bond=5):
|
||||
bs, n, n_atom = X.shape
|
||||
node_indices = X.argmax(-1) # (bs, n)
|
||||
|
||||
exp_ind1 = node_indices[ :, :, None, None, None].expand(bs, n, n_atom, n_bond, n_bond)
|
||||
exp_ind2 = node_indices[ :, :, None, None, None].expand(bs, n, n, n_bond, n_bond)
|
||||
|
||||
q_e = torch.gather(q_e, 1, exp_ind1)
|
||||
q_e = torch.gather(q_e, 2, exp_ind2) # (bs, n, n, n_bond, n_bond)
|
||||
|
||||
|
||||
node_mask = X.sum(-1) != 0
|
||||
no_edge = (~node_mask)[:, :, None] & (~node_mask)[:, None, :]
|
||||
q_e[no_edge] = torch.tensor([1, 0, 0, 0, 0]).type_as(q_e)
|
||||
|
||||
return q_e
|
30
graph_dit/diffusion/distributions.py
Normal file
30
graph_dit/diffusion/distributions.py
Normal file
@@ -0,0 +1,30 @@
|
||||
import torch
|
||||
|
||||
class DistributionNodes:
|
||||
def __init__(self, histogram):
|
||||
""" Compute the distribution of the number of nodes in the dataset, and sample from this distribution.
|
||||
historgram: dict. The keys are num_nodes, the values are counts
|
||||
"""
|
||||
|
||||
if type(histogram) == dict:
|
||||
max_n_nodes = max(histogram.keys())
|
||||
prob = torch.zeros(max_n_nodes + 1)
|
||||
for num_nodes, count in histogram.items():
|
||||
prob[num_nodes] = count
|
||||
else:
|
||||
prob = histogram
|
||||
|
||||
self.prob = prob / prob.sum()
|
||||
self.m = torch.distributions.Categorical(prob)
|
||||
|
||||
def sample_n(self, n_samples, device):
|
||||
idx = self.m.sample((n_samples,))
|
||||
return idx.to(device)
|
||||
|
||||
def log_prob(self, batch_n_nodes):
|
||||
assert len(batch_n_nodes.size()) == 1
|
||||
p = self.prob.to(batch_n_nodes.device)
|
||||
|
||||
probas = p[batch_n_nodes]
|
||||
log_p = torch.log(probas + 1e-30)
|
||||
return log_p
|
159
graph_dit/diffusion/noise_schedule.py
Normal file
159
graph_dit/diffusion/noise_schedule.py
Normal file
@@ -0,0 +1,159 @@
|
||||
import torch
|
||||
import utils
|
||||
from diffusion import diffusion_utils
|
||||
|
||||
class PredefinedNoiseScheduleDiscrete(torch.nn.Module):
|
||||
def __init__(self, noise_schedule, timesteps):
|
||||
super(PredefinedNoiseScheduleDiscrete, self).__init__()
|
||||
self.timesteps = timesteps
|
||||
|
||||
if noise_schedule == 'cosine':
|
||||
betas = diffusion_utils.cosine_beta_schedule_discrete(timesteps)
|
||||
elif noise_schedule == 'custom':
|
||||
betas = diffusion_utils.custom_beta_schedule_discrete(timesteps)
|
||||
else:
|
||||
raise NotImplementedError(noise_schedule)
|
||||
|
||||
self.register_buffer('betas', torch.from_numpy(betas).float())
|
||||
|
||||
# 0.9999
|
||||
self.alphas = 1 - torch.clamp(self.betas, min=0, max=1)
|
||||
|
||||
log_alpha = torch.log(self.alphas)
|
||||
log_alpha_bar = torch.cumsum(log_alpha, dim=0)
|
||||
self.alphas_bar = torch.exp(log_alpha_bar)
|
||||
|
||||
def forward(self, t_normalized=None, t_int=None):
|
||||
assert int(t_normalized is None) + int(t_int is None) == 1
|
||||
if t_int is None:
|
||||
t_int = torch.round(t_normalized * self.timesteps)
|
||||
return self.betas[t_int.long()]
|
||||
|
||||
def get_alpha_bar(self, t_normalized=None, t_int=None):
|
||||
assert int(t_normalized is None) + int(t_int is None) == 1
|
||||
if t_int is None:
|
||||
t_int = torch.round(t_normalized * self.timesteps)
|
||||
### new
|
||||
self.alphas_bar = self.alphas_bar.to(t_int.device)
|
||||
return self.alphas_bar[t_int.long()]
|
||||
|
||||
|
||||
class DiscreteUniformTransition:
|
||||
def __init__(self, x_classes: int, e_classes: int, y_classes: int):
|
||||
self.X_classes = x_classes
|
||||
self.E_classes = e_classes
|
||||
self.y_classes = y_classes
|
||||
self.u_x = torch.ones(1, self.X_classes, self.X_classes)
|
||||
if self.X_classes > 0:
|
||||
self.u_x = self.u_x / self.X_classes
|
||||
|
||||
self.u_e = torch.ones(1, self.E_classes, self.E_classes)
|
||||
if self.E_classes > 0:
|
||||
self.u_e = self.u_e / self.E_classes
|
||||
|
||||
self.u_y = torch.ones(1, self.y_classes, self.y_classes)
|
||||
if self.y_classes > 0:
|
||||
self.u_y = self.u_y / self.y_classes
|
||||
|
||||
def get_Qt(self, beta_t, device, X=None, flatten_e=None):
|
||||
""" Returns one-step transition matrices for X and E, from step t - 1 to step t.
|
||||
Qt = (1 - beta_t) * I + beta_t / K
|
||||
|
||||
beta_t: (bs) noise level between 0 and 1
|
||||
returns: qx (bs, dx, dx), qe (bs, de, de), qy (bs, dy, dy).
|
||||
"""
|
||||
beta_t = beta_t.unsqueeze(1)
|
||||
beta_t = beta_t.to(device)
|
||||
self.u_x = self.u_x.to(device)
|
||||
self.u_e = self.u_e.to(device)
|
||||
self.u_y = self.u_y.to(device)
|
||||
|
||||
q_x = beta_t * self.u_x + (1 - beta_t) * torch.eye(self.X_classes, device=device).unsqueeze(0)
|
||||
q_e = beta_t * self.u_e + (1 - beta_t) * torch.eye(self.E_classes, device=device).unsqueeze(0)
|
||||
q_y = beta_t * self.u_y + (1 - beta_t) * torch.eye(self.y_classes, device=device).unsqueeze(0)
|
||||
|
||||
return utils.PlaceHolder(X=q_x, E=q_e, y=q_y)
|
||||
|
||||
def get_Qt_bar(self, alpha_bar_t, device, X=None, flatten_e=None):
|
||||
""" Returns t-step transition matrices for X and E, from step 0 to step t.
|
||||
Qt = prod(1 - beta_t) * I + (1 - prod(1 - beta_t)) / K
|
||||
|
||||
alpha_bar_t: (bs) Product of the (1 - beta_t) for each time step from 0 to t.
|
||||
returns: qx (bs, dx, dx), qe (bs, de, de), qy (bs, dy, dy).
|
||||
"""
|
||||
alpha_bar_t = alpha_bar_t.unsqueeze(1)
|
||||
alpha_bar_t = alpha_bar_t.to(device)
|
||||
self.u_x = self.u_x.to(device)
|
||||
self.u_e = self.u_e.to(device)
|
||||
self.u_y = self.u_y.to(device)
|
||||
|
||||
q_x = alpha_bar_t * torch.eye(self.X_classes, device=device).unsqueeze(0) + (1 - alpha_bar_t) * self.u_x
|
||||
q_e = alpha_bar_t * torch.eye(self.E_classes, device=device).unsqueeze(0) + (1 - alpha_bar_t) * self.u_e
|
||||
q_y = alpha_bar_t * torch.eye(self.y_classes, device=device).unsqueeze(0) + (1 - alpha_bar_t) * self.u_y
|
||||
|
||||
return utils.PlaceHolder(X=q_x, E=q_e, y=q_y)
|
||||
|
||||
|
||||
class MarginalTransition:
|
||||
def __init__(self, x_marginals, e_marginals, xe_conditions, ex_conditions, y_classes, n_nodes):
|
||||
self.X_classes = len(x_marginals)
|
||||
self.E_classes = len(e_marginals)
|
||||
self.y_classes = y_classes
|
||||
self.x_marginals = x_marginals # Dx
|
||||
self.e_marginals = e_marginals # Dx, De
|
||||
self.xe_conditions = xe_conditions
|
||||
|
||||
self.u_x = x_marginals.unsqueeze(0).expand(self.X_classes, -1).unsqueeze(0) # 1, Dx, Dx
|
||||
self.u_e = e_marginals.unsqueeze(0).expand(self.E_classes, -1).unsqueeze(0) # 1, De, De
|
||||
self.u_xe = xe_conditions.unsqueeze(0) # 1, Dx, De
|
||||
self.u_ex = ex_conditions.unsqueeze(0) # 1, De, Dx
|
||||
self.u = self.get_union_transition(self.u_x, self.u_e, self.u_xe, self.u_ex, n_nodes) # 1, Dx + n*De, Dx + n*De
|
||||
|
||||
def get_union_transition(self, u_x, u_e, u_xe, u_ex, n_nodes):
|
||||
u_e = u_e.repeat(1, n_nodes, n_nodes) # (1, n*de, n*de)
|
||||
u_xe = u_xe.repeat(1, 1, n_nodes) # (1, dx, n*de)
|
||||
u_ex = u_ex.repeat(1, n_nodes, 1) # (1, n*de, dx)
|
||||
u0 = torch.cat([u_x, u_xe], dim=2) # (1, dx, dx + n*de)
|
||||
u1 = torch.cat([u_ex, u_e], dim=2) # (1, n*de, dx + n*de)
|
||||
u = torch.cat([u0, u1], dim=1) # (1, dx + n*de, dx + n*de)
|
||||
return u
|
||||
|
||||
def index_edge_margin(self, X, q_e, n_bond=5):
|
||||
# q_e: (bs, dx, de) --> (bs, n, de)
|
||||
bs, n, n_atom = X.shape
|
||||
node_indices = X.argmax(-1) # (bs, n)
|
||||
ind = node_indices[ :, :, None].expand(bs, n, n_bond)
|
||||
q_e = torch.gather(q_e, 1, ind)
|
||||
return q_e
|
||||
|
||||
def get_Qt(self, beta_t, device):
|
||||
""" Returns one-step transition matrices for X and E, from step t - 1 to step t.
|
||||
Qt = (1 - beta_t) * I + beta_t / K
|
||||
beta_t: (bs)
|
||||
returns: q (bs, d0, d0)
|
||||
"""
|
||||
bs = beta_t.size(0)
|
||||
d0 = self.u.size(-1)
|
||||
self.u = self.u.to(device)
|
||||
u = self.u.expand(bs, d0, d0)
|
||||
|
||||
beta_t = beta_t.to(device)
|
||||
beta_t = beta_t.view(bs, 1, 1)
|
||||
q = beta_t * u + (1 - beta_t) * torch.eye(d0, device=device).unsqueeze(0)
|
||||
|
||||
return utils.PlaceHolder(X=q, E=None, y=None)
|
||||
|
||||
def get_Qt_bar(self, alpha_bar_t, device):
|
||||
""" Returns t-step transition matrices for X and E, from step 0 to step t.
|
||||
Qt = prod(1 - beta_t) * I + (1 - prod(1 - beta_t)) * K
|
||||
alpha_bar_t: (bs, 1) roduct of the (1 - beta_t) for each time step from 0 to t.
|
||||
returns: q (bs, d0, d0)
|
||||
"""
|
||||
bs = alpha_bar_t.size(0)
|
||||
d0 = self.u.size(-1)
|
||||
alpha_bar_t = alpha_bar_t.to(device)
|
||||
alpha_bar_t = alpha_bar_t.view(bs, 1, 1)
|
||||
self.u = self.u.to(device)
|
||||
q = alpha_bar_t * torch.eye(d0, device=device).unsqueeze(0) + (1 - alpha_bar_t) * self.u
|
||||
|
||||
return utils.PlaceHolder(X=q, E=None, y=None)
|
Reference in New Issue
Block a user