File size: 2,784 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 | import torch
import torch.nn.functional as F
def warmup_lr(step):
return min(step, 5000) / 5000
def training_with_esm(x_0_con, x_0_dis, trainer, trainer_dis, ns_con, ns_dis, trans, FLAGS, cls_head, class_weights):
t = torch.randint(FLAGS.T, size=(x_0_con['input_ids'].shape[0], ), device=x_0_con['input_ids'].device)
pt = torch.ones_like(t).float() / FLAGS.T
embeddings = trainer.model.get_embeddings(x_0_con)
noise = torch.randn_like(embeddings)
x_t_con = trainer.make_x_t(embeddings, t, noise)
torch.cuda.empty_cache()
log_x_start = torch.log(x_0_dis.float().clamp(min=1e-30))
x_t_dis = trainer_dis.q_sample(log_x_start=log_x_start, t=t)
eps = trainer.model(x_t_con, t, x_t_dis.to(x_t_con.device))
ps_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps)
con_loss = F.mse_loss(eps, noise, reduction='none')
con_loss = con_loss.mean()
kl, ps_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, x_t_con)
ps_0_dis = torch.exp(ps_0_dis)
kl_prior = trainer_dis.kl_prior(log_x_start)
dis_loss = (kl / pt + kl_prior).mean()
torch.cuda.empty_cache()
noise_ns = torch.randn_like(ns_con)
ns_t_con = trainer.make_x_t(ns_con, t, noise_ns)
log_ns_start = torch.log(ns_dis.float().clamp(min=1e-30))
ns_t_dis = trainer_dis.q_sample(log_x_start=log_ns_start, t=t)
torch.cuda.empty_cache()
eps_ns = trainer.model(x_t_con, t, ns_t_dis.to(ns_t_dis.device))
ns_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps_ns)
_, ns_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, ns_t_con)
ns_0_dis = torch.exp(ns_0_dis)
torch.cuda.empty_cache()
triplet_loss = torch.nn.TripletMarginLoss(margin=0.5, p=2)
triplet_con = triplet_loss(embeddings, ps_0_con, ns_0_con)
total_loss = con_loss
st=0
triplet_dis = []
for item in trans.output_info:
ed = st + item[0]
ps_dis = F.cross_entropy(ps_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none')
ns_dis = F.cross_entropy(ns_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none')
triplet_dis.append(max((ps_dis-ns_dis).mean()+1,0))
st = ed
triplet_dis = sum(triplet_dis)/len(triplet_dis)
return total_loss, triplet_con, dis_loss, triplet_dis
def make_negative_condition_esm(x_0_con, x_0_dis, base_model):
device = x_0_dis.device
with torch.no_grad():
_, _, embeddings, _ = base_model(
input_ids=x_0_con['input_ids'],
attention_masks=x_0_con['attention_mask']
)
B = embeddings.shape[0]
indices = torch.randperm(B)
ns_con = embeddings[indices].clone()
ns_dis = x_0_dis[indices].clone()
return ns_con, ns_dis
|