import torch from torch.utils.tensorboard import SummaryWriter from torch.utils.data import DataLoader import numpy as np from sklearn.metrics import * from omegaconf import OmegaConf import os import random from mcts import MCTS import esm from encoders import AptaBLE from utils import get_scores, API_Dataset, get_nt_esm_dataset, rna2vec from accelerate import Accelerator import glob import os import requests from transformers import AutoTokenizer, AutoModelForMaskedLM # accelerator = Accelerator(kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]) # NOTE: Buggy | Disables unused parameter issue accelerator = Accelerator() class AptaBLE_Pipeline(): """In-house API prediction score pipeline.""" def __init__(self, lr, dropout, weight_decay, epochs, model_type, model_version, model_save_path, accelerate_save_path, tensorboard_logdir, *args, **kwargs): self.device = accelerator.device self.lr = lr self.weight_decay = weight_decay self.epochs = epochs self.model_type = model_type self.model_version = model_version self.model_save_path = model_save_path self.accelerate_save_path = accelerate_save_path self.tensorboard_logdir = tensorboard_logdir esm_prot_encoder, self.esm_alphabet = esm.pretrained.esm.pretrained.esm2_t33_650M_UR50D() # ESM-2 Encoder # Freeze ESM-2 for name, param in esm_prot_encoder.named_parameters(): param.requires_grad = False for name, param in esm_prot_encoder.named_parameters(): if "layers.30" in name or "layers.31" in name or "layers.32" in name: param.requires_grad = True self.batch_converter = self.esm_alphabet.get_batch_converter(truncation_seq_length=1678) # self.nt_tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-2.5b-1000g") # nt_encoder = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-2.5b-1000g") self.nt_tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-v2-50m-multi-species", trust_remote_code=True) nt_encoder = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-v2-50m-multi-species", trust_remote_code=True) self.model = AptaBLE( apta_encoder=nt_encoder, prot_encoder=esm_prot_encoder, dropout=dropout, ).to(self.device) self.criterion = torch.nn.BCELoss().to(self.device) def train(self): print('Training the model!') # Initialize writer instance writer = SummaryWriter(log_dir=f"log/{self.model_type}/{self.model_version}") # Initialize early stopping self.early_stopper = EarlyStopper(3, 3) self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay) self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, [4, 7, 10], 0.1) # Configure pytorch objects for distributed environment (i.e. sharded dataloader, multiple copies of model, etc.) self.model, self.optimizer, self.train_loader, self.test_loader, self.bench_loader, self.scheduler = accelerator.prepare(self.model, self.optimizer, self.train_loader, self.test_loader, self.bench_loader, self.scheduler) best_loss = 100 for epoch in range(1, self.epochs+1): self.model.train() loss_train, _, _ = self.batch_step(self.train_loader, train_mode=True) self.model.eval() self.scheduler.step() with torch.no_grad(): loss_test, pred_test, target_test = self.batch_step(self.test_loader, train_mode=False) test_scores = get_scores(target_test, pred_test) print("\tTrain Loss: {: .6f}\tTest Loss: {: .6f}\tTest ACC: {:.6f}\tTest AUC: {:.6f}\tTest MCC: {:.6f}\tTest PR_AUC: {:.6f}\tF1: {:.6f}\n".format(loss_train ,loss_test, test_scores['acc'], test_scores['roc_auc'], test_scores['mcc'], test_scores['pr_auc'], test_scores['f1'])) # stop_early = self.early_stopper.early_stop(loss_test) # Early stop - model has not improved on eval set. # if stop_early: # break # Only do checkpointing after near-convergence if epoch > 2: with torch.no_grad(): loss_bench, pred_bench, target_bench = self.batch_step(self.bench_loader, train_mode=False) bench_scores = get_scores(target_bench, pred_bench) print("\Bench Loss: {: .6f}\Bench ACC: {:.6f}\Bench AUC: {:.6f}\tBench MCC: {:.6f}\tBench PR_AUC: {:.6f}\tBench F1: {:.6f}\n".format(loss_bench, bench_scores['acc'], bench_scores['roc_auc'], bench_scores['mcc'], bench_scores['pr_auc'], bench_scores['f1'])) writer.add_scalar("Loss/bench", loss_bench, epoch) for k, v in bench_scores.items(): if isinstance(v, float): writer.add_scalar(f'{k}/bench', bench_scores[k], epoch) # Checkpoint based off of benchmark criteria # If model has improved and early stopping patience counter was just reset: if bench_scores['mcc'] > 0.5 and test_scores['mcc'] > 0.5 and loss_bench < 0.9 and accelerator.is_main_process: best_loss = loss_test # Remove all other files # for f in glob.glob(f'{self.model_save_path}/model*.pt'): # os.remove(f) accelerator.save_state(self.accelerate_save_path) model = accelerator.unwrap_model(self.model) torch.save(model.state_dict(), f'{self.model_save_path}/model_epoch={epoch}.pt') print(f'Model saved at {self.model_save_path}') print(f'Accelerate statistics saved at {self.accelerate_save_path}!') # Access via accelerator.load_state("./output") # logging writer.add_scalar("Loss/train", loss_train, epoch) writer.add_scalar("Loss/test", loss_test, epoch) for k, v in test_scores.items(): if isinstance(v, float): writer.add_scalar(f'{k}/test', test_scores[k], epoch) print("Training finished | access tensorboard via 'tensorboard --logdir=runs'.") writer.flush() writer.close() def batch_step(self, loader, train_mode = True): loss_total = 0 pred = np.array([]) target = np.array([]) for batch_idx, (apta, esm_prot, y, apta_attn, prot_attn) in enumerate(loader): if train_mode: self.optimizer.zero_grad() y_pred = self.predict(apta, esm_prot, apta_attn, prot_attn) y_true = torch.tensor(y, dtype=torch.float32).to(self.device) # not needed since accelerator modifies dataloader to automatically map input objects to correct dev loss = self.criterion(torch.flatten(y_pred), y_true) if train_mode: accelerator.backward(loss) # Accelerate backward() method scales gradients and uses appropriate backward method as configured across devices self.optimizer.step() loss_total += loss.item() pred = np.append(pred, torch.flatten(y_pred).clone().detach().cpu().numpy()) target = np.append(target, torch.flatten(y_true).clone().detach().cpu().numpy()) mode = 'train' if train_mode else 'eval' print(mode + "[{}/{}({:.0f}%)]".format(batch_idx, len(loader), 100. * batch_idx / len(loader)), end = "\r", flush=True) loss_total /= len(loader) return loss_total, pred, target def predict(self, apta, esm_prot, apta_attn, prot_attn): y_pred, _, _, _ = self.model(apta, esm_prot, apta_attn, prot_attn) return y_pred def inference(self, apta, prot, labels): """Perform inference on a batch of aptamer/protein pairs.""" self.model.eval() max_length = 275#nt_tokenizer.model_max_length inputs = [(i, j) for i, j in zip(labels, prot)] _, _, prot_tokens = self.batch_converter(inputs) apta_toks = self.nt_tokenizer.batch_encode_plus(apta, return_tensors='pt', padding='max_length', max_length=max_length)['input_ids'] apta_attention_mask = apta_toks != self.nt_tokenizer.pad_token_id # # truncating prot_tokenized = prot_tokens[:, :1680] # # padding prot_ex = torch.ones((prot_tokenized.shape[0], 1680), dtype=torch.int64)*self.esm_alphabet.padding_idx prot_ex[:, :prot_tokenized.shape[1]] = prot_tokenized prot_attention_mask = prot_ex != self.esm_alphabet.padding_idx loader = DataLoader(API_Dataset(apta_toks, prot_ex, labels, apta_attention_mask, prot_attention_mask), batch_size=1, shuffle=False) self.model, loader = accelerator.prepare(self.model, loader) with torch.no_grad(): _, pred, _ = self.batch_step(loader, train_mode=False) return pred def recommend(self, target, n_aptamers, depth, iteration, verbose=True): candidates = [] _, _, prot_tokens = self.batch_converter([(1, target)]) prot_tokenized = torch.tensor(prot_tokens, dtype=torch.int64) # adjusting for max protein sequence length during model training encoded_targetprotein = torch.ones((prot_tokenized.shape[0], 1678), dtype=torch.int64)*self.esm_alphabet.padding_idx encoded_targetprotein[:, :prot_tokenized.shape[1]] = prot_tokenized encoded_targetprotein = encoded_targetprotein.to(self.device) mcts = MCTS(encoded_targetprotein, depth=depth, iteration=iteration, states=8, target_protein=target, device=self.device, esm_alphabet=self.esm_alphabet) for _ in range(n_aptamers): mcts.make_candidate(self.model) candidates.append(mcts.get_candidate()) self.model.eval() with torch.no_grad(): sim_seq = np.array([mcts.get_candidate()]) print('first candidate: ', sim_seq) # apta = torch.tensor(rna2vec(sim_seq), dtype=torch.int64).to(self.device) apta = self.nt_tokenizer.batch_encode_plus(sim_seq, return_tensors='pt', padding='max_length', max_length=275)['input_ids'] apta_attn = apta != self.nt_tokenizer.pad_token_id prot_attn = encoded_targetprotein != self.esm_alphabet.padding_idx score, _, _, _ = self.model(apta.to(self.device), encoded_targetprotein.to(self.device), apta_attn.to(self.device), prot_attn.to(self.device)) if verbose: candidate = mcts.get_candidate() print("candidate:\t", candidate, "\tscore:\t", score) print("*"*80) mcts.reset() def set_data_for_training(self, filepath, batch_size): # ds_train, ds_test, ds_bench = get_nt_esm_dataset(filepath, self.nt_tokenizer, self.batch_converter, self.esm_alphabet) ds_train, ds_test, ds_bench = get_nt_esm_dataset(filepath, self.nt_tokenizer, self.batch_converter, self.esm_alphabet) self.train_loader = DataLoader(API_Dataset(ds_train[0], ds_train[1], ds_train[2], ds_train[3], ds_train[4]), batch_size=batch_size, shuffle=True) self.test_loader = DataLoader(API_Dataset(ds_test[0], ds_test[1], ds_test[2], ds_test[3], ds_test[4]), batch_size=batch_size, shuffle=False) self.bench_loader = DataLoader(API_Dataset(ds_bench[0], ds_bench[1], ds_bench[2], ds_bench[3], ds_bench[4]), batch_size=batch_size, shuffle=False) class EarlyStopper: def __init__(self, patience=1, min_delta=0): self.patience = patience self.min_delta = min_delta self.counter = 0 self.min_validation_loss = float('inf') def early_stop(self, validation_loss): if validation_loss < self.min_validation_loss: self.min_validation_loss = validation_loss self.counter = 0 elif validation_loss > (self.min_validation_loss + self.min_delta): self.counter += 1 if self.counter >= self.patience: return True return False def seed_torch(seed=5471): random.seed(seed) os.environ['PYTHONHASHSEED'] = str(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) # if you are using multi-GPU. torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True def main(): conf = OmegaConf.load('config.yaml') hyperparameters = conf.hyperparameters logging = conf.logging lr = hyperparameters['lr'] wd = hyperparameters['weight_decay'] dropout = hyperparameters['dropout'] batch_size = hyperparameters['batch_size'] epochs = hyperparameters['epochs'] model_type = logging['model_type'] model_version = logging['model_version'] model_save_path = logging['model_save_path'] accelerate_save_path = logging['accelerate_save_path'] tensorboard_logdir = logging['tensorboard_logdir'] seed = hyperparameters['seed'] if not os.path.exists(model_save_path): os.makedirs(model_save_path) seed_torch(seed=seed) pipeline = AptaBLE_Pipeline( lr=lr, weight_decay=wd, epochs=epochs, model_type=model_type, model_version=model_version, model_save_path=model_save_path, accelerate_save_path=accelerate_save_path, tensorboard_logdir=tensorboard_logdir, d_model=128, d_ff=512, n_layers=6, n_heads=8, dropout=dropout, load_best_pt=True, # already loads the pretrained model using the datasets included in repo -- no need to run the bottom two cells device='cuda', seed=seed) datapath = "./data/ABW_real_dna_aptamers_HC_v6.pkl" # datapath = './data/ABW_real_dna_aptamers_HC_neg_scrambles_neg_homology.pkl' pipeline.set_data_for_training(datapath, batch_size=batch_size) pipeline.train() endpoint = 'https://slack.atombioworks.com/hooks/t3y99qu6pi81frhwrhef1849wh' msg = {"text": "Model has finished training."} _ = requests.post(endpoint, json=msg, headers={"Content-Type": "application/json"}, ) return if __name__ == "__main__": # launch training w/ the following: "accelerate launch api_prediction.py [args]" main()