| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| import math |
| from .layers import get_timestep_embedding |
|
|
| def extract(v, t, x_shape): |
| out = torch.gather(v, index=t, dim=0).float() |
| return out.view([t.shape[0]] + [1] * (len(x_shape) - 1)) |
|
|
|
|
| class GaussianDiffusionTrainer(nn.Module): |
| def __init__(self, model, beta_1, beta_T, T): |
| super().__init__() |
|
|
| self.model = model |
| self.T = T |
| betas = torch.linspace(beta_1, beta_T, T, dtype=torch.float64).double() |
| alphas = 1. - betas |
| self.register_buffer('betas', betas) |
| alphas_bar = torch.cumprod(alphas, dim=0) |
|
|
| self.register_buffer( |
| 'sqrt_alphas_bar', torch.sqrt(alphas_bar)) |
| self.register_buffer( |
| 'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar)) |
| self.register_buffer( |
| 'sqrt_recip_alphas_bar', torch.sqrt(1. / alphas_bar)) |
| self.register_buffer( |
| 'sqrt_recipm1_alphas_bar', torch.sqrt(1. / alphas_bar - 1)) |
|
|
| def make_x_t(self, x_0_con, t, noise): |
| x_t_con = ( |
| extract(self.sqrt_alphas_bar, t, x_0_con.shape) * x_0_con + |
| extract(self.sqrt_one_minus_alphas_bar, t, x_0_con.shape) * noise) |
| return x_t_con |
|
|
| def predict_xstart_from_eps(self, x_t, t, eps): |
| assert x_t.shape == eps.shape |
| return ( |
| extract(self.sqrt_recip_alphas_bar, t, x_t.shape) * x_t - |
| extract(self.sqrt_recipm1_alphas_bar, t, x_t.shape) * eps |
| ) |
|
|
|
|
| class GaussianDiffusionSampler(nn.Module): |
| def __init__(self, model, beta_1, beta_T, T, |
| mean_type='eps', var_type='fixedlarge'): |
| assert mean_type in ['xprev' 'xstart', 'epsilon'] |
| assert var_type in ['fixedlarge', 'fixedsmall'] |
| super().__init__() |
|
|
| self.model = model |
| self.T = T |
| self.mean_type = mean_type |
| self.var_type = var_type |
|
|
| betas = torch.linspace(beta_1, beta_T, T, dtype=torch.float64).double() |
|
|
| alphas = 1. - betas |
| self.register_buffer( |
| 'betas', betas) |
| alphas_bar = torch.cumprod(alphas, dim=0) |
| alphas_bar_prev = F.pad(alphas_bar, [1, 0], value=1)[:T] |
|
|
| self.register_buffer( |
| 'sqrt_alphas_bar', torch.sqrt(alphas_bar)) |
| self.register_buffer( |
| 'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar)) |
|
|
| self.register_buffer( |
| 'sqrt_recip_alphas_bar', torch.sqrt(1. / alphas_bar)) |
| self.register_buffer( |
| 'sqrt_recipm1_alphas_bar', torch.sqrt(1. / alphas_bar - 1)) |
|
|
| self.register_buffer( |
| 'posterior_var', |
| self.betas * (1. - alphas_bar_prev) / (1. - alphas_bar)) |
| |
| self.register_buffer( |
| 'posterior_log_var_clipped', |
| torch.log( |
| torch.cat([self.posterior_var[1:2], self.posterior_var[1:]]))) |
| self.register_buffer( |
| 'posterior_mean_coef1', |
| torch.sqrt(alphas_bar_prev) * self.betas / (1. - alphas_bar)) |
| self.register_buffer( |
| 'posterior_mean_coef2', |
| torch.sqrt(alphas) * (1. - alphas_bar_prev) / (1. - alphas_bar)) |
|
|
| def q_mean_variance(self, x_0, x_t, t): |
| """Mean and log variance of q(x_{t-1} | x_t, x_0).""" |
| assert x_0.shape == x_t.shape |
| posterior_mean = ( |
| extract(self.posterior_mean_coef1, t, x_t.shape) * x_0 + |
| extract(self.posterior_mean_coef2, t, x_t.shape) * x_t |
| ) |
| posterior_log_var_clipped = extract( |
| self.posterior_log_var_clipped, t, x_t.shape) |
| return posterior_mean, posterior_log_var_clipped |
|
|
| def predict_xstart_from_eps(self, x_t, t, eps): |
| assert x_t.shape == eps.shape |
| return ( |
| extract(self.sqrt_recip_alphas_bar, t, x_t.shape) * x_t - |
| extract(self.sqrt_recipm1_alphas_bar, t, x_t.shape) * eps |
| ) |
|
|
|
|
| def p_mean_variance(self, x_t, t, cond, trans): |
| model_log_var = { |
| 'fixedlarge': torch.log(torch.cat([self.posterior_var[1:2], |
| self.betas[1:]])), |
| 'fixedsmall': self.posterior_log_var_clipped, |
| }[self.var_type] |
| model_log_var = extract(model_log_var, t, x_t.shape) |
|
|
| if self.mean_type == 'epsilon': |
| eps = self.model(x_t, t, cond) |
| x_0 = self.predict_xstart_from_eps(x_t, t, eps=eps) |
| model_mean, _ = self.q_mean_variance(x_0, x_t, t) |
| else: |
| raise NotImplementedError(self.mean_type) |
|
|
| return model_mean, model_log_var |
|
|
|
|
| class DiffusionProteinModel(nn.Module): |
| def __init__(self, base_model, FLAGS): |
| super().__init__() |
| self.base_model = base_model |
| self._freeze_base_model_layers() |
| self.embed_dim = FLAGS.nf |
|
|
| self.time_embed = nn.Sequential( |
| nn.Linear(self.embed_dim, self.embed_dim * 4), |
| nn.SiLU(), |
| nn.Linear(self.embed_dim * 4, 1280) |
| ) |
|
|
| self.cond_proj = nn.Linear(FLAGS.cond_size, 1280) |
|
|
| self.output_proj = nn.Sequential( |
| nn.Linear(1280, FLAGS.output_size * 2), |
| nn.SiLU(), |
| nn.Linear(FLAGS.output_size * 2, FLAGS.output_size) |
| ) |
| def _freeze_base_model_layers(self): |
| for param in self.base_model.parameters(): |
| param.requires_grad = False |
|
|
| for i in range(23, 33): |
| for param in self.base_model.esm_model.encoder.layer[i].parameters(): |
| param.requires_grad = True |
|
|
| for component in [self.base_model.lm_head, self.base_model.attn_head, |
| self.base_model.attn_ln, self.base_model.attn_skip, |
| self.base_model.linear_layers, self.base_model.reduction_layers, |
| self.base_model.clf_ln]: |
| for param in component.parameters(): |
| param.requires_grad = True |
|
|
| for param in self.base_model.esm_model.embeddings.parameters(): |
| param.requires_grad = False |
| def decode_embeddings(self, embeddings): |
| mlm_logits = self.base_model.lm_head(embeddings) |
| return mlm_logits |
|
|
| def get_embeddings(self, x, detach=False): |
| _, _, x_combined, _ = self.base_model( |
| input_ids=x['input_ids'], |
| attention_masks=x['attention_mask'] |
| ) |
| if detach: |
| return x_combined.detach() |
| return x_combined |
|
|
| def forward(self, x, time_cond, cond): |
| if isinstance(x, dict): |
| _, _, x_combined, _ = self.base_model( |
| input_ids=x['input_ids'], |
| attention_masks=x['attention_mask'] |
| ) |
| x = x_combined |
| seq_length = x.size(1) |
|
|
| temb = get_timestep_embedding(time_cond, self.embed_dim) |
| temb = self.time_embed(temb) |
| temb = temb.unsqueeze(1).expand(-1, seq_length, -1) |
|
|
| if len(cond.shape) == 3: |
| cond = self.cond_proj(cond) |
| else: |
| cond = self.cond_proj(cond) |
| cond = cond.unsqueeze(1).expand(-1, seq_length, -1) |
|
|
| x = x + temb + cond |
| return x |
|
|
| def get_token_embeddings(self): |
| return self.base_model.esm_model.embeddings.word_embeddings.weight |
|
|
