Spaces:
Running on Zero
Running on Zero
| import os | |
| import time | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import pytorch_lightning as pl | |
| import wandb | |
| from src.models.transformer_model import GraphTransformer | |
| from src.diffusion.noise_schedule import PredefinedNoiseSchedule | |
| from src.diffusion import diffusion_utils | |
| from src.metrics.train_metrics import TrainLoss | |
| from src.metrics.abstract_metrics import SumExceptBatchMetric, SumExceptBatchMSE, NLL | |
| import src.utils | |
| class LiftedDenoisingDiffusion(pl.LightningModule): | |
| def __init__(self, cfg, dataset_infos, train_metrics, sampling_metrics, visualization_tools, extra_features=None, | |
| domain_features=None): | |
| super().__init__() | |
| input_dims = dataset_infos.input_dims | |
| output_dims = dataset_infos.output_dims | |
| nodes_dist = dataset_infos.nodes_dist | |
| self.norm_values = cfg.model.normalize_factors | |
| self.norm_biases = cfg.model.norm_biases | |
| self.gamma = PredefinedNoiseSchedule(cfg.model.diffusion_noise_schedule, timesteps=cfg.model.diffusion_steps) | |
| diffusion_utils.check_issues_norm_values(self.gamma, self.norm_values[1], self.norm_values[2]) | |
| self.cfg = cfg | |
| self.name = cfg.general.name | |
| self.model_dtype = torch.float32 | |
| self.T = cfg.model.diffusion_steps | |
| self.Xdim = input_dims['X'] | |
| self.Edim = input_dims['E'] | |
| self.ydim = input_dims['y'] | |
| self.Xdim_output = output_dims['X'] | |
| self.Edim_output = output_dims['E'] | |
| self.ydim_output = output_dims['y'] | |
| self.node_dist = nodes_dist | |
| self.dataset_info = dataset_infos | |
| self.val_nll = NLL() | |
| self.val_X_mse = SumExceptBatchMSE() | |
| self.val_E_mse = SumExceptBatchMSE() | |
| self.val_y_mse = SumExceptBatchMSE() | |
| self.val_X_logp = SumExceptBatchMetric() | |
| self.val_E_logp = SumExceptBatchMetric() | |
| self.val_y_logp = SumExceptBatchMSE() | |
| self.test_nll = NLL() | |
| self.test_X_mse = SumExceptBatchMSE() | |
| self.test_E_mse = SumExceptBatchMSE() | |
| self.test_y_mse = SumExceptBatchMSE() | |
| self.test_X_logp = SumExceptBatchMetric() | |
| self.test_E_logp = SumExceptBatchMetric() | |
| self.test_y_logp = SumExceptBatchMSE() | |
| self.train_loss = TrainLoss() | |
| self.train_metrics = train_metrics | |
| self.sampling_metrics = sampling_metrics | |
| self.visualization_tools = visualization_tools | |
| self.save_hyperparameters(ignore=['train_metrics', 'sampling_metrics']) | |
| self.visualization_tools = visualization_tools | |
| self.model = GraphTransformer(n_layers=cfg.model.n_layers, | |
| input_dims=input_dims, | |
| hidden_mlp_dims=cfg.model.hidden_mlp_dims, | |
| hidden_dims=cfg.model.hidden_dims, | |
| output_dims=output_dims, | |
| act_fn_in=nn.ReLU(), | |
| act_fn_out=nn.ReLU()) | |
| self.save_hyperparameters() | |
| self.start_epoch_time = None | |
| self.train_iterations = None | |
| self.val_iterations = None | |
| self.log_every_steps = cfg.general.log_every_steps | |
| self.number_chain_steps = cfg.general.number_chain_steps | |
| self.best_val_nll = 1e8 | |
| self.val_counter = 0 | |
| def training_step(self, data, i): | |
| dense_data, node_mask = utils.to_dense(x=data.x, edge_index=data.edge_index, edge_attr=data.edge_attr, | |
| batch=data.batch) | |
| dense_data = dense_data.mask(node_mask) | |
| X, E = dense_data.X, dense_data.E | |
| normalized_data = utils.normalize(X, E, data.y, self.norm_values, self.norm_biases, node_mask) | |
| noisy_data = self.apply_noise(normalized_data.X, normalized_data.E, normalized_data.y, node_mask) | |
| extra_data = self.compute_extra_data(noisy_data) | |
| pred = self.forward(noisy_data, extra_data, node_mask) | |
| # TODO: change noisy data | |
| mse = self.train_loss(masked_pred_epsX=pred.X, | |
| masked_pred_epsE=pred.E, | |
| pred_y=pred.y, | |
| true_epsX=noisy_data['epsX'], | |
| true_epsE=noisy_data['epsE'], | |
| true_y=noisy_data['epsy'], | |
| log=i % self.log_every_steps == 0) | |
| self.train_metrics(masked_pred_epsX=pred.X, | |
| masked_pred_epsE=pred.E, | |
| pred_y=pred.y, | |
| true_epsX=noisy_data['epsX'], | |
| true_epsE=noisy_data['epsE'], | |
| true_y=noisy_data['epsy'], log=i % self.log_every_steps == 0) | |
| return {'loss': mse} | |
| def configure_optimizers(self): | |
| return torch.optim.AdamW(self.parameters(), lr=self.cfg.train.lr, amsgrad=True, | |
| weight_decay=self.cfg.train.weight_decay) | |
| def on_fit_start(self) -> None: | |
| self.train_iterations = len(self.trainer.datamodule.train_dataloader()) | |
| if self.local_rank == 0: | |
| utils.setup_wandb(self.cfg) | |
| def on_train_epoch_start(self) -> None: | |
| self.start_epoch_time = time.time() | |
| self.train_loss.reset() | |
| self.train_metrics.reset() | |
| def on_train_epoch_end(self) -> None: | |
| to_log = self.train_loss.log_epoch_metrics() | |
| self.print(f"Epoch {self.current_epoch}: X_mse: {to_log['train_epoch/epoch_X_mse'] :.3f}" | |
| f" -- E mse: {to_log['train_epoch/epoch_E_mse'] :.3f} --" | |
| f" y_mse: {to_log['train_epoch/epoch_y_mse'] :.3f}" | |
| f" -- {time.time() - self.start_epoch_time:.1f}s ") | |
| epoch_at_metrics, epoch_bond_metrics = self.train_metrics.log_epoch_metrics() | |
| self.print(f"Epoch {self.current_epoch}: {epoch_at_metrics} -- {epoch_bond_metrics}") | |
| def on_validation_epoch_start(self) -> None: | |
| self.val_nll.reset() | |
| self.val_X_mse.reset() | |
| self.val_E_mse.reset() | |
| self.val_y_mse.reset() | |
| self.val_X_logp.reset() | |
| self.val_E_logp.reset() | |
| self.val_y_logp.reset() | |
| def validation_step(self, data, i): | |
| dense_data, node_mask = utils.to_dense(x=data.x, edge_index=data.edge_index, edge_attr=data.edge_attr, | |
| batch=data.batch) | |
| dense_data = dense_data.mask(node_mask) | |
| X, E = dense_data.X, dense_data.E | |
| normalized_data = utils.normalize(X, E, data.y, self.norm_values, self.norm_biases, node_mask) | |
| noisy_data = self.apply_noise(normalized_data.X, normalized_data.E, data.y, node_mask) | |
| extra_data = self.compute_extra_data(noisy_data) | |
| pred = self.forward(noisy_data, extra_data, node_mask) | |
| # TODO: check if compute val loss should be called on the normalized data or not | |
| nll = self.compute_val_loss(pred, noisy_data, normalized_data.X, normalized_data.E, normalized_data.y, | |
| node_mask, test=False) | |
| return {'loss': nll} | |
| def on_validation_epoch_end(self) -> None: | |
| metrics = [self.val_nll.compute(), self.val_X_mse.compute(), self.val_E_mse.compute(), | |
| self.val_y_mse.compute(), self.val_X_logp.compute(), self.val_E_logp.compute(), | |
| self.val_y_logp.compute()] | |
| if wandb.run: | |
| wandb.log({"val/epoch_NLL": metrics[0], | |
| "val/X_mse": metrics[1], | |
| "val/E_mse": metrics[2], | |
| "val/y_mse": metrics[3], | |
| "val/X_logp": metrics[4], | |
| "val/E_logp": metrics[5], | |
| "val/y_logp": metrics[6]}, commit=False) | |
| print(f"Epoch {self.current_epoch}: Val NLL {metrics[0] :.2f} -- Val Atom type MSE {metrics[1] :.2f} -- ", | |
| f"Val Edge type MSE: {metrics[2] :.2f} -- Val Global feat. MSE {metrics[3] :.2f}", | |
| f"-- Val X Reconstruction loss {metrics[4] :.2f} -- Val E Reconstruction loss {metrics[5] :.2f}", | |
| f"-- Val y Reconstruction loss {metrics[6] : .2f}\n") | |
| # Log val nll with default Lightning logger, so it can be monitored by checkpoint callback | |
| val_nll = metrics[0] | |
| self.log("val/epoch_NLL", val_nll, sync_dist=True) | |
| if wandb.run: | |
| wandb.log(self.log_info(), commit=False) | |
| if val_nll < self.best_val_nll: | |
| self.best_val_nll = val_nll | |
| print('Val loss: %.4f \t Best val loss: %.4f\n' % (val_nll, self.best_val_nll)) | |
| self.val_counter += 1 | |
| if self.val_counter % self.cfg.general.sample_every_val == 0: | |
| start = time.time() | |
| samples_left_to_generate = self.cfg.general.samples_to_generate | |
| samples_left_to_save = self.cfg.general.samples_to_save | |
| chains_left_to_save = self.cfg.general.chains_to_save | |
| samples = [] | |
| ident = 0 | |
| while samples_left_to_generate > 0: | |
| bs = 2 * self.cfg.train.batch_size | |
| to_generate = min(samples_left_to_generate, bs) | |
| to_save = min(samples_left_to_save, bs) | |
| chains_save = min(chains_left_to_save, bs) | |
| samples.extend(self.sample_batch(batch_id=ident, | |
| batch_size=to_generate, | |
| num_nodes=None, save_final=to_save, | |
| keep_chain=chains_save, | |
| number_chain_steps=self.number_chain_steps)) | |
| ident += to_generate | |
| samples_left_to_save -= to_save | |
| samples_left_to_generate -= to_generate | |
| chains_left_to_save -= chains_save | |
| self.sampling_metrics(samples, self.name, self.current_epoch, val_counter=-1, test=False) | |
| print(f'Sampling took {time.time() - start:.2f} seconds\n') | |
| self.sampling_metrics.reset() | |
| def on_test_epoch_start(self) -> None: | |
| self.test_nll.reset() | |
| self.test_X_mse.reset() | |
| self.test_E_mse.reset() | |
| self.test_y_mse.reset() | |
| self.test_X_logp.reset() | |
| self.test_E_logp.reset() | |
| self.test_y_logp.reset() | |
| if self.local_rank == 0: | |
| utils.setup_wandb(self.cfg) | |
| def test_step(self, data, i): | |
| dense_data, node_mask = utils.to_dense(x=data.x, edge_index=data.edge_index, edge_attr=data.edge_attr, | |
| batch=data.batch) | |
| dense_data = dense_data.mask(node_mask) | |
| X, E = dense_data.X, dense_data.E | |
| normalized_data = utils.normalize(X, E, data.y, self.norm_values, self.norm_biases, node_mask) | |
| noisy_data = self.apply_noise(normalized_data.X, normalized_data.E, normalized_data.y, node_mask) | |
| extra_data = self.compute_extra_data(noisy_data) | |
| pred = self.forward(noisy_data, extra_data, node_mask) | |
| nll = self.compute_val_loss(pred, noisy_data, normalized_data.X, normalized_data.E, | |
| normalized_data.y, node_mask, test=True) | |
| return {'loss': nll} | |
| def on_test_epoch_end(self) -> None: | |
| """ Measure likelihood on a test set and compute stability metrics. """ | |
| metrics = [self.test_nll.compute(), self.test_X_mse.compute(), self.test_E_mse.compute(), | |
| self.test_y_mse.compute(), self.test_X_logp.compute(), self.test_E_logp.compute(), | |
| self.test_y_logp.compute()] | |
| log_dict={"test/epoch_NLL": metrics[0], | |
| "test/X_mse": metrics[1], | |
| "test/E_mse": metrics[2], | |
| "test/y_mse": metrics[3], | |
| "test/X_logp": metrics[4], | |
| "test/E_logp": metrics[5], | |
| "test/y_logp": metrics[6]} | |
| if wandb.run: | |
| wandb.log(log_dict, commit=False) | |
| print(f"Epoch {self.current_epoch}: Test NLL {metrics[0] :.2f} -- Test Atom type MSE {metrics[1] :.2f} -- ", | |
| f"Test Edge type MSE: {metrics[2] :.2f} -- Test Global feat. MSE {metrics[3] :.2f}", | |
| f"-- Test X Reconstruction loss {metrics[4] :.2f} -- Test E Reconstruction loss {metrics[5] :.2f}", | |
| f"-- Test y Reconstruction loss {metrics[6] : .2f}\n") | |
| test_nll = metrics[0] | |
| if wandb.run: | |
| wandb.log({"test/epoch_NLL": test_nll}, commit=False) | |
| wandb.log(self.log_info(), commit=False) | |
| print(f'Test loss: {test_nll :.4f}') | |
| samples_left_to_generate = self.cfg.general.final_model_samples_to_generate | |
| samples_left_to_save = self.cfg.general.final_model_samples_to_save | |
| chains_left_to_save = self.cfg.general.final_model_chains_to_save | |
| samples = [] | |
| id = 0 | |
| while samples_left_to_generate > 0: | |
| bs = 2 * self.cfg.train.batch_size | |
| to_generate = min(samples_left_to_generate, bs) | |
| to_save = min(samples_left_to_save, bs) | |
| chains_save = min(chains_left_to_save, bs) | |
| samples.extend(self.sample_batch(id, to_generate, num_nodes=None, save_final=to_save, | |
| keep_chain=chains_save, number_chain_steps=self.number_chain_steps)) | |
| id += to_generate | |
| samples_left_to_save -= to_save | |
| samples_left_to_generate -= to_generate | |
| chains_left_to_save -= chains_save | |
| self.sampling_metrics.reset() | |
| self.sampling_metrics(samples, self.name, self.current_epoch, self.val_counter, test=True) | |
| self.sampling_metrics.reset() | |
| def kl_prior(self, X, E, y, node_mask): | |
| """Computes the KL between q(z1 | x) and the prior p(z1) = Normal(0, 1). | |
| This is essentially a lot of work for something that is in practice negligible in the loss. However, you | |
| compute it so that you see it when you've made a mistake in your noise schedule. | |
| """ | |
| # Compute the last alpha value, alpha_T. | |
| ones = torch.ones((X.size(0), 1)) | |
| ones = ones.type_as(X) | |
| gamma_T = self.gamma(ones) | |
| alpha_T = diffusion_utils.alpha(gamma_T, X.size()) | |
| # Compute means. | |
| mu_T_X = alpha_T * X | |
| mu_T_E = alpha_T.unsqueeze(1) * E | |
| mu_T_y = alpha_T.squeeze(1) * y | |
| # Compute standard deviations (only batch axis for x-part, inflated for h-part). | |
| sigma_T_X = diffusion_utils.sigma(gamma_T, mu_T_X.size()) | |
| sigma_T_E = diffusion_utils.sigma(gamma_T, mu_T_E.size()) | |
| sigma_T_y = diffusion_utils.sigma(gamma_T, mu_T_y.size()) | |
| # Compute KL for h-part. | |
| kl_distance_X = diffusion_utils.gaussian_KL(mu_T_X, sigma_T_X) | |
| kl_distance_E = diffusion_utils.gaussian_KL(mu_T_E, sigma_T_E) | |
| kl_distance_y = diffusion_utils.gaussian_KL(mu_T_y, sigma_T_y) | |
| return kl_distance_X + kl_distance_E + kl_distance_y | |
| def log_constants_p_y_given_z0(self, batch_size): | |
| """ Computes p(y|z0)= -0.5 ydim (log(2pi) + gamma(0)). | |
| sigma_y = sqrt(sigma_0^2 / alpha_0^2) = SNR(-0.5 gamma_0). | |
| output size: (batch_size) | |
| """ | |
| if self.ydim_output == 0: | |
| return 0.0 | |
| zeros = torch.zeros((batch_size, 1)) | |
| gamma_0 = self.gamma(zeros).squeeze(1) | |
| # Recall that | |
| return -0.5 * self.ydim * (gamma_0 + np.log(2 * np.pi)) | |
| def reconstruction_logp(self, data, data_0, gamma_0, eps, pred_0, node_mask, epsilon=1e-10, test=False): | |
| """ Reconstruction loss. | |
| output size: (1). | |
| """ | |
| X, E, y = data.values() | |
| X_0, E_0, y_0 = data_0.values() | |
| # TODO: why don't we need the values of X and E? | |
| _, _, eps_y0 = eps.values() | |
| predy = pred_0.y | |
| # 1. Compute reconstruction loss for global, continuous features | |
| if test: | |
| error_y = -0.5 * self.test_y_logp(predy, eps_y0) | |
| else: | |
| error_y = -0.5 * self.val_y_logp(predy, eps_y0) | |
| # The _constants_ depending on sigma_0 from the cross entropy term E_q(z0 | y) [log p(y | z0)]. | |
| neg_log_constants = - self.log_constants_p_y_given_z0(y.shape[0]) | |
| log_py = error_y + neg_log_constants | |
| # 2. Compute reconstruction loss for integer/categorical features on nodes and edges | |
| # Compute sigma_0 and rescale to the integer scale of the data_utils. | |
| sigma_0 = diffusion_utils.sigma(gamma_0, target_shape=X_0.size()) | |
| sigma_0_X = sigma_0 * self.norm_values[0] | |
| sigma_0_E = (sigma_0 * self.norm_values[1]).unsqueeze(-1) | |
| # Unnormalize features | |
| unnormalized_data = utils.unnormalize(X, E, y, self.norm_values, self.norm_biases, node_mask, collapse=False) | |
| unnormalized_0 = utils.unnormalize(X_0, E_0, y_0, self.norm_values, self.norm_biases, node_mask, collapse=False) | |
| X_0, E_0, _ = unnormalized_0.X, unnormalized_0.E, unnormalized_0.y | |
| assert unnormalized_data.X.size() == X_0.size() | |
| # Centered cat features around 1, since onehot encoded. | |
| E_0_centered = E_0 - 1 | |
| X_0_centered = X_0 - 1 | |
| # Compute integrals from 0.5 to 1.5 of the normal distribution | |
| log_pE_proportional = torch.log( | |
| diffusion_utils.cdf_std_gaussian((E_0_centered + 0.5) / sigma_0_E) | |
| - diffusion_utils.cdf_std_gaussian((E_0_centered - 0.5) / sigma_0_E) | |
| + epsilon) | |
| log_pX_proportional = torch.log( | |
| diffusion_utils.cdf_std_gaussian((X_0_centered + 0.5) / sigma_0_X) | |
| - diffusion_utils.cdf_std_gaussian((X_0_centered - 0.5) / sigma_0_X) | |
| + epsilon) | |
| # Normalize the distributions over the categories. | |
| norm_cst_E = torch.logsumexp(log_pE_proportional, dim=-1, keepdim=True) | |
| norm_cst_X = torch.logsumexp(log_pX_proportional, dim=-1, keepdim=True) | |
| log_probabilities_E = log_pE_proportional - norm_cst_E | |
| log_probabilities_X = log_pX_proportional - norm_cst_X | |
| # Select the log_prob of the current category using the one-hot representation. | |
| logps = utils.PlaceHolder(X=log_probabilities_X * unnormalized_data.X, | |
| E=log_probabilities_E * unnormalized_data.E, | |
| y=None).mask(node_mask) | |
| if test: | |
| log_pE = - self.test_E_logp(-logps.E) | |
| log_pX = - self.test_X_logp(-logps.X) | |
| else: | |
| log_pE = - self.val_E_logp(-logps.E) | |
| log_pX = - self.val_X_logp(-logps.X) | |
| return log_pE + log_pX + log_py | |
| def apply_noise(self, X, E, y, node_mask): | |
| """ Sample noise and apply it to the data. """ | |
| # When evaluating, the loss for t=0 is computed separately | |
| lowest_t = 0 if self.training else 1 | |
| # Sample a timestep t. | |
| t_int = torch.randint(lowest_t, self.T + 1, size=(X.size(0), 1)) | |
| t_int = t_int.type_as(X).float() # (bs, 1) | |
| s_int = t_int - 1 | |
| # Normalize t to [0, 1]. Note that the negative | |
| # step of s will never be used, since then p(x | z0) is computed. | |
| s_normalized = s_int / self.T | |
| t_normalized = t_int / self.T | |
| # Compute gamma_s and gamma_t via the network. | |
| gamma_s = diffusion_utils.inflate_batch_array(self.gamma(s_normalized), X.size()) # (bs, 1, 1), | |
| gamma_t = diffusion_utils.inflate_batch_array(self.gamma(t_normalized), X.size()) # (bs, 1, 1) | |
| # Compute alpha_t and sigma_t from gamma, with correct size for X, E and z | |
| alpha_t = diffusion_utils.alpha(gamma_t, X.size()) # (bs, 1, ..., 1), same n_dims than X | |
| sigma_t = diffusion_utils.sigma(gamma_t, X.size()) # (bs, 1, ..., 1), same n_dims than X | |
| # Sample zt ~ Normal(alpha_t x, sigma_t) | |
| eps = diffusion_utils.sample_feature_noise(X.size(), E.size(), y.size(), node_mask).type_as(X) | |
| # Sample z_t given x, h for timestep t, from q(z_t | x, h) | |
| X_t = alpha_t * X + sigma_t * eps.X | |
| E_t = alpha_t.unsqueeze(1) * E + sigma_t.unsqueeze(1) * eps.E | |
| y_t = alpha_t.squeeze(1) * y + sigma_t.squeeze(1) * eps.y | |
| noisy_data = {'t': t_normalized, 's': s_normalized, 'gamma_t': gamma_t, 'gamma_s': gamma_s, | |
| 'epsX': eps.X, 'epsE': eps.E, 'epsy': eps.y, | |
| 'X_t': X_t, 'E_t': E_t, 'y_t': y_t, 't_int': t_int} | |
| return noisy_data | |
| def compute_val_loss(self, pred, noisy_data, X, E, y, node_mask, test=False): | |
| """ Computes an estimator for the variational lower bound, or the simple loss (MSE). | |
| pred: (batch_size, n, total_features) | |
| noisy_data: dict | |
| X, E, y : (bs, n, dx), (bs, n, n, de), (bs, dy) | |
| node_mask : (bs, n) | |
| Output: nll (size 1). """ | |
| s = noisy_data['s'] | |
| gamma_s = noisy_data['gamma_s'] # gamma_s.size() == X.size() | |
| gamma_t = noisy_data['gamma_t'] | |
| epsX = noisy_data['epsX'] | |
| epsE = noisy_data['epsE'] | |
| epsy = noisy_data['epsy'] | |
| X_t = noisy_data['X_t'] | |
| E_t = noisy_data['E_t'] | |
| y_t = noisy_data['y_t'] | |
| # 1. | |
| N = node_mask.sum(1).long() | |
| log_pN = self.node_dist.log_prob(N) | |
| # 2. The KL between q(z_T | x) and p(z_T) = Normal(0, 1). Should be close to zero. Do not forget the prefactor | |
| kl_prior_without_prefactor = self.kl_prior(X, E, y, node_mask) | |
| delta_log_py = -self.ydim_output * np.log(self.norm_values[2]) | |
| delta_log_px = -self.Xdim_output * N * np.log(self.norm_values[0]) | |
| delta_log_pE = -self.Edim_output * 0.5 * N * (N-1) * np.log(self.norm_values[1]) | |
| kl_prior = kl_prior_without_prefactor - delta_log_px - delta_log_py - delta_log_pE | |
| # 3. Diffusion loss | |
| # Compute weighting with SNR: (1 - SNR(s-t)) for epsilon parametrization. | |
| SNR_weight = - (1 - diffusion_utils.SNR(gamma_s - gamma_t)) | |
| sqrt_SNR_weight = torch.sqrt(SNR_weight) # same n_dims than X | |
| # Compute the error. | |
| weighted_predX_diffusion = sqrt_SNR_weight * pred.X | |
| weighted_epsX_diffusion = sqrt_SNR_weight * epsX | |
| weighted_predE_diffusion = sqrt_SNR_weight.unsqueeze(1) * pred.E | |
| weighted_epsE_diffusion = sqrt_SNR_weight.unsqueeze(1) * epsE | |
| weighted_predy_diffusion = sqrt_SNR_weight.squeeze(1) * pred.y | |
| weighted_epsy_diffusion = sqrt_SNR_weight.squeeze(1) * epsy | |
| # Compute the MSE summed over channels | |
| if test: | |
| diffusion_error = (self.test_X_mse(weighted_predX_diffusion, weighted_epsX_diffusion) + | |
| self.test_E_mse(weighted_predE_diffusion, weighted_epsE_diffusion) + | |
| self.test_y_mse(weighted_predy_diffusion, weighted_epsy_diffusion)) | |
| else: | |
| diffusion_error = (self.val_X_mse(weighted_predX_diffusion, weighted_epsX_diffusion) + | |
| self.val_E_mse(weighted_predE_diffusion, weighted_epsE_diffusion) + | |
| self.val_y_mse(weighted_predy_diffusion, weighted_epsy_diffusion)) | |
| loss_all_t = 0.5 * self.T * diffusion_error # t=0 is not included here. | |
| # 4. Compute L0 term : -log p (X, E, y | z_0) = reconstruction loss | |
| # Compute noise values for t = 0. | |
| t_zeros = torch.zeros_like(s) # bs, 1 | |
| gamma_0 = diffusion_utils.inflate_batch_array(self.gamma(t_zeros), X_t.size()) # bs, 1, 1 | |
| alpha_0 = diffusion_utils.alpha(gamma_0, X_t.size()) # bs, 1, 1 | |
| sigma_0 = diffusion_utils.sigma(gamma_0, X_t.size()) # bs, 1, 1 | |
| # Sample z_0 given X, E, y for timestep t, from q(z_t | X, E, y) | |
| eps0 = diffusion_utils.sample_feature_noise(X_t.size(), E_t.size(), y_t.size(), node_mask).type_as(X_t) | |
| X_0 = alpha_0 * X_t + sigma_0 * eps0.X | |
| E_0 = alpha_0.unsqueeze(1) * E_t + sigma_0.unsqueeze(1) * eps0.E | |
| y_0 = alpha_0.squeeze(1) * y_t + sigma_0.squeeze(1) * eps0.y | |
| noisy_data0 = {'X_t': X_0, 'E_t': E_0, 'y_t': y_0, 't': t_zeros} | |
| extra_data = self.compute_extra_data(noisy_data) | |
| pred_0 = self.forward(noisy_data0, extra_data, node_mask) | |
| loss_term_0 = - self.reconstruction_logp(data={'X': X, 'E': E, 'y': y}, | |
| data_0={'X_0': X_0, 'E_0': E_0, 'y_0': y_0}, | |
| gamma_0=gamma_0, | |
| eps={'eps_X0': eps0.X, 'eps_E0': eps0.E, 'eps_y0': eps0.y}, | |
| pred_0=pred_0, | |
| node_mask=node_mask, | |
| test=test) | |
| # Combine terms | |
| nlls = - log_pN + kl_prior + loss_all_t + loss_term_0 | |
| assert len(nlls.shape) == 1, f'{nlls.shape} has more than only batch dim.' | |
| # Update NLL metric object and return batch nll | |
| nll = self.test_nll(nlls) if test else self.val_nll(nlls) # Average over the batch | |
| wandb.log({"kl prior": kl_prior.mean(), | |
| "Estimator loss terms": loss_all_t.mean(), | |
| "Loss term 0": loss_term_0, | |
| "log_pn": log_pN.mean(), | |
| 'test_nll' if test else 'val_nll': nll}, | |
| commit=False) | |
| return nll | |
| def forward(self, noisy_data, extra_data, node_mask): | |
| """ Concatenates extra data to the noisy data, then calls the network. """ | |
| X = torch.cat((noisy_data['X_t'], extra_data.X), dim=2) | |
| E = torch.cat((noisy_data['E_t'], extra_data.E), dim=3) | |
| y = torch.hstack((noisy_data['y_t'], extra_data.y)) | |
| return self.model(X, E, y, node_mask) | |
| def log_info(self): | |
| """ | |
| Some info logging of the model. | |
| """ | |
| gamma_0 = self.gamma(torch.zeros(1, device=self.device)) | |
| gamma_1 = self.gamma(torch.ones(1, device=self.device)) | |
| log_SNR_max = -gamma_0 | |
| log_SNR_min = -gamma_1 | |
| info = {'log_SNR_max': log_SNR_max.item(), 'log_SNR_min': log_SNR_min.item()} | |
| print("", info, "\n") | |
| return info | |
| def sample_batch(self, batch_id: int, batch_size: int, keep_chain: int, save_final: int, number_chain_steps: int, | |
| num_nodes=None): | |
| """ | |
| :param batch_id: int | |
| :param batch_size: int | |
| :param num_nodes: int, <int>tensor (batch_size) (optional) for specifying number of nodes | |
| :param save_final: int: number of predictions to save to file | |
| :param keep_chain: int: number of chains to save to file | |
| :param number_chain_steps: number of timesteps to save for each chain | |
| :return: molecule_list. Each element of this list is a tuple (atom_types, charges, positions) | |
| """ | |
| if num_nodes is None: | |
| n_nodes = self.node_dist.sample_n(batch_size, self.device) | |
| elif type(num_nodes) == int: | |
| n_nodes = num_nodes * torch.ones(batch_size, device=self.device, dtype=torch.int) | |
| else: | |
| assert isinstance(num_nodes, torch.Tensor) | |
| n_nodes = num_nodes | |
| n_nodes_max = torch.max(n_nodes).item() | |
| # Build the masks | |
| arange = torch.arange(n_nodes_max, device=self.device).unsqueeze(0).expand(batch_size, -1) | |
| node_mask = arange < n_nodes.unsqueeze(1) | |
| node_mask = node_mask.float() | |
| # Sample noise -- z has size (n_samples, n_nodes, n_features) | |
| # TODO: how to move on the right device in the multi-gpu case? | |
| z_T = diffusion_utils.sample_feature_noise(X_size=(batch_size, n_nodes_max, self.Xdim_output), | |
| E_size=(batch_size, n_nodes_max, n_nodes_max, self.Edim_output), | |
| y_size=(batch_size, self.ydim_output), | |
| node_mask=node_mask) | |
| X, E, y = z_T.X, z_T.E, z_T.y | |
| assert (E == torch.transpose(E, 1, 2)).all() | |
| assert number_chain_steps < self.T | |
| chain_X_size = torch.Size((number_chain_steps, keep_chain, X.size(1))) | |
| chain_E_size = torch.Size((number_chain_steps, keep_chain, E.size(1), E.size(2))) | |
| chain_X = torch.zeros(chain_X_size) | |
| chain_E = torch.zeros(chain_E_size) | |
| # Iteratively sample p(z_s | z_t) for t = 1, ..., T, with s = t - 1. | |
| average_X_coord = [] | |
| average_E_coord = [] | |
| for s_int in reversed(range(0, self.T)): | |
| s_array = s_int * torch.ones((batch_size, 1)).type_as(y) | |
| t_array = s_array + 1 | |
| s_norm = s_array / self.T | |
| t_norm = t_array / self.T | |
| z_s = self.sample_p_zs_given_zt(s=s_norm, t=t_norm, X_t=X, E_t=E, y_t=y, node_mask=node_mask) | |
| X, E, y = z_s.X, z_s.E, z_s.y | |
| write_index = (s_int * number_chain_steps) // self.T | |
| unnormalized = utils.unnormalize(X=X[:keep_chain], E=E[:keep_chain], y=y[:keep_chain], | |
| norm_values=self.norm_values, | |
| norm_biases=self.norm_biases, | |
| node_mask=node_mask[:keep_chain], | |
| collapse=True) | |
| chain_X[write_index] = unnormalized.X | |
| chain_E[write_index] = unnormalized.E | |
| average_X_coord.append(X.abs().mean().item()) | |
| average_E_coord.append(E.abs().mean().item()) | |
| print(f"Average X coordinate at each step {[int(c) for i, c in enumerate(average_X_coord) if i % 10 == 0]}") | |
| print(f"Average E coordinate at each step {[int(c) for i, c in enumerate(average_E_coord) if i % 10 == 0]}") | |
| # Finally sample the discrete data given the last latent code z0 | |
| final_graph = self.sample_discrete_graph_given_z0(X, E, y, node_mask) | |
| X, E, y = final_graph.X, final_graph.E, final_graph.y | |
| assert (E == torch.transpose(E, 1, 2)).all() | |
| print("Examples of generated graphs:") | |
| for i in range(min(5, X.shape[0])): | |
| print("E", E[i]) | |
| print("X: ", X[i]) | |
| # Prepare the chain for saving | |
| if keep_chain > 0: | |
| final_X_chain = X[:keep_chain] | |
| final_E_chain = E[:keep_chain] | |
| chain_X[0] = final_X_chain | |
| chain_E[0] = final_E_chain | |
| chain_X = diffusion_utils.reverse_tensor(chain_X) | |
| chain_E = diffusion_utils.reverse_tensor(chain_E) | |
| # Repeat last frame to see final sample better | |
| chain_X = torch.cat([chain_X, chain_X[-1:].repeat(10, 1, 1)], dim=0) | |
| chain_E = torch.cat([chain_E, chain_E[-1:].repeat(10, 1, 1, 1)], dim=0) | |
| assert chain_X.size(0) == (number_chain_steps + 10) | |
| # Split the generated molecules | |
| molecule_list = [] | |
| for i in range(batch_size): | |
| n = n_nodes[i] | |
| atom_types = X[i, :n].cpu() | |
| edge_types = E[i, :n, :n].cpu() | |
| molecule_list.append([atom_types, edge_types]) | |
| # Visualize chains | |
| if self.visualization_tools is not None: | |
| print('Visualizing chains...') | |
| current_path = os.getcwd() | |
| num_molecules = chain_X.size(1) # number of molecules | |
| for i in range(num_molecules): | |
| result_path = os.path.join(current_path, f'chains/{self.cfg.general.name}/' | |
| f'epoch{self.current_epoch}/' | |
| f'chains/molecule_{batch_id + i}') | |
| if not os.path.exists(result_path): | |
| os.makedirs(result_path) | |
| _ = self.visualization_tools.visualize_chain(result_path, | |
| chain_X[:, i, :].numpy(), | |
| chain_E[:, i, :].numpy()) | |
| print('\r{}/{} complete'.format(i+1, num_molecules), end='', flush=True) | |
| # Visualize the final molecules | |
| print("Visualizing molecules...") | |
| current_path = os.getcwd() | |
| result_path = os.path.join(current_path, | |
| f'graphs/{self.name}/epoch{self.current_epoch}_b{batch_id}/') | |
| self.visualization_tools.visualize(result_path, molecule_list, save_final, log='graph') | |
| print("Done.") | |
| return molecule_list | |
| def sample_discrete_graph_given_z0(self, X_0, E_0, y_0, node_mask): | |
| """ Samples X, E, y ~ p(X, E, y|z0): once the diffusion is done, we need to map the result | |
| to categorical values. | |
| """ | |
| zeros = torch.zeros(size=(X_0.size(0), 1), device=X_0.device) | |
| gamma_0 = self.gamma(zeros) | |
| # Computes sqrt(sigma_0^2 / alpha_0^2) | |
| sigma = diffusion_utils.SNR(-0.5 * gamma_0).unsqueeze(1) | |
| noisy_data = {'X_t': X_0, 'E_t': E_0, 'y_t': y_0, 't': torch.zeros(y_0.shape[0], 1).type_as(y_0)} | |
| extra_data = self.compute_extra_data(noisy_data) | |
| eps0 = self.forward(noisy_data, extra_data, node_mask) | |
| # Compute mu for p(zs | zt). | |
| sigma_0 = diffusion_utils.sigma(gamma_0, target_shape=eps0.X.size()) | |
| alpha_0 = diffusion_utils.alpha(gamma_0, target_shape=eps0.X.size()) | |
| pred_X = 1. / alpha_0 * (X_0 - sigma_0 * eps0.X) | |
| pred_E = 1. / alpha_0.unsqueeze(1) * (E_0 - sigma_0.unsqueeze(1) * eps0.E) | |
| pred_y = 1. / alpha_0.squeeze(1) * (y_0 - sigma_0.squeeze(1) * eps0.y) | |
| assert (pred_E == torch.transpose(pred_E, 1, 2)).all() | |
| sampled = diffusion_utils.sample_normal(pred_X, pred_E, pred_y, sigma, node_mask).type_as(pred_X) | |
| assert (sampled.E == torch.transpose(sampled.E, 1, 2)).all() | |
| sampled = utils.unnormalize(sampled.X, sampled.E, sampled.y, self.norm_values, | |
| self.norm_biases, node_mask, collapse=True) | |
| return sampled | |
| def sample_p_zs_given_zt(self, s, t, X_t, E_t, y_t, node_mask): | |
| """Samples from zs ~ p(zs | zt). Only used during sampling.""" | |
| gamma_s = self.gamma(s) | |
| gamma_t = self.gamma(t) | |
| sigma2_t_given_s, sigma_t_given_s, alpha_t_given_s = diffusion_utils.sigma_and_alpha_t_given_s(gamma_t, | |
| gamma_s, | |
| X_t.size()) | |
| sigma_s = diffusion_utils.sigma(gamma_s, target_shape=X_t.size()) | |
| sigma_t = diffusion_utils.sigma(gamma_t, target_shape=X_t.size()) | |
| E_t = (E_t + E_t.transpose(1, 2)) / 2 | |
| noisy_data = {'X_t': X_t, 'E_t': E_t, 'y_t': y_t, 't': t} | |
| extra_data = self.compute_extra_data(noisy_data) | |
| eps = self.forward(noisy_data, extra_data, node_mask) | |
| # Compute mu for p(zs | zt). | |
| mu_X = X_t / alpha_t_given_s - (sigma2_t_given_s / (alpha_t_given_s * sigma_t)) * eps.X | |
| mu_E = E_t / alpha_t_given_s.unsqueeze(1) - (sigma2_t_given_s / (alpha_t_given_s * sigma_t)).unsqueeze(1) * eps.E | |
| mu_y = y_t / alpha_t_given_s.squeeze(1) - (sigma2_t_given_s / (alpha_t_given_s * sigma_t)).squeeze(1) * eps.y | |
| # Compute sigma for p(zs | zt). | |
| sigma = sigma_t_given_s * sigma_s / sigma_t | |
| # Sample zs given the parameters derived from zt. | |
| z_s = diffusion_utils.sample_normal(mu_X, mu_E, mu_y, sigma, node_mask).type_as(mu_X) | |
| return z_s | |
| def compute_extra_data(self, noisy_data): | |
| """ At every training step (after adding noise) and step in sampling, compute extra information and append to | |
| the network input. """ | |
| X = noisy_data['X_t'] | |
| E = noisy_data['E_t'] | |
| extra_x = torch.zeros((*X.shape[:-1], 0)).type_as(X) | |
| extra_edge_attr = torch.zeros((*E.shape[:-1], 0)).type_as(E) | |
| t = noisy_data['t'] | |
| return utils.PlaceHolder(X=extra_x, E=extra_edge_attr, y=t) | |