File size: 7,321 Bytes
aa45001 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 | 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))
# Avoid log(0) at the first diffusion step.
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
|