File size: 9,739 Bytes
d8ed92a |
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 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 |
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from transformers import AutoModel, AutoTokenizer
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
from tqdm import tqdm
from datetime import datetime
import pandas as pd
import numpy as np
import pickle
import os
# Hyperparameters dictionary
path = "/workspace/sg666/MDpLM"
hyperparams = {
"batch_size": 1,
"learning_rate": 5e-4,
"num_epochs": 5,
"esm_model_path": "facebook/esm2_t33_650M_UR50D",
'mlm_model_path': path + "/benchmarks/MLM/model_ckpts/best_model_epoch",
"mdlm_model_path": path + "/checkpoints/membrane_automodel/epochs30_lr3e-4_bsz16_gradclip1_beta-one0.9_beta-two0.999_bf16_all-params",
"train_data": path + "/benchmarks/Supervised/Localization/true_deeploc2.0_cell-local_train-val.csv",
"test_data" : path + "/benchmarks/Supervised/Localization/true_deeploc2.0_cell-local_test.csv",
}
# Helper functions to obtain all embeddings for a sequence
def load_models(esm_model_path, mlm_model_path, mdlm_model_path):
esm_tokenizer = AutoTokenizer.from_pretrained(esm_model_path)
esm_model = AutoModel.from_pretrained(esm_model_path).to(device)
mlm_model = AutoModel.from_pretrained(mlm_model_path).to(device)
mdlm_model = AutoModel.from_pretrained(mdlm_model_path).to(device)
return esm_tokenizer, esm_model, mlm_model, mdlm_model
def get_latents(embedding_type, tokenizer, esm_model, mlm_model, mdlm_model, sequence, device):
if embedding_type == "esm":
inputs = tokenizer(sequence, return_tensors='pt').to(device)
with torch.no_grad():
embeddings = esm_model(**inputs).last_hidden_state.squeeze(0)
elif embedding_type == "mlm":
inputs = tokenizer(sequence, return_tensors='pt')['input_ids'].to(device)
with torch.no_grad():
embeddings = mlm_model(inputs).last_hidden_state.squeeze(0)
elif embedding_type == "mdlm":
inputs = tokenizer(sequence, return_tensors='pt')['input_ids'].to(device)
with torch.no_grad():
embeddings = mdlm_model(inputs).last_hidden_state.squeeze(0)
return embeddings
# Dataset class can load pickle file
class LocalizationDataset(Dataset):
def __init__(self, embedding_type, csv_file, esm_model_path, mlm_model_path, mdlm_model_path, device):
self.data = pd.read_csv(csv_file)
self.data = self.data[self.data['Sequence'].apply(len) < 1024].reset_index(drop=True)
self.embedding_type = embedding_type
self.tokenizer, self.esm_model, self.mlm_model, self.mdlm_model = load_models(esm_model_path, mlm_model_path, mdlm_model_path)
self.device = device
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
sequence = self.data.iloc[idx]['Sequence']
embeddings = get_latents(self.embedding_type, self.tokenizer, self.mlm_model, self.esm_model, self.mdlm_model,
sequence, self.device)
label = 0 if self.data.iloc[idx]['Cell membrane'] == 0 else 1
labels = torch.tensor(label, dtype=torch.float32).view(1,1).squeeze(-1)
return embeddings, labels
# Predict localization with MLP head using pooled embeddings
class LocalizationPredictor(nn.Module):
def __init__(self, input_dim):
super(LocalizationPredictor, self).__init__()
self.classifier = nn.Sequential(
nn.Linear(input_dim, 640),
nn.ReLU(),
nn.Linear(640, 1)
)
def forward(self, embeddings):
logits = self.classifier(embeddings)
logits = torch.mean(logits, dim=1)
probs = torch.nn.functional.softmax(logits)
return probs
# Training function
def train(model, dataloader, optimizer, criterion, device):
model.train()
total_loss = 0
for embeddings, labels in tqdm(dataloader):
embeddings, labels = embeddings.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(embeddings)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / len(dataloader)
# Evaluation function
def evaluate(model, dataloader, device):
model.eval()
preds, true_labels = [], []
with torch.no_grad():
for embeddings, labels in tqdm(dataloader):
embeddings, labels = embeddings.to(device), labels.to(device)
outputs = model(embeddings)
preds.append(outputs.cpu().numpy())
true_labels.append(labels.cpu().numpy())
return preds, true_labels
# Metrics calculation
def calculate_metrics(preds, labels, threshold=0.5):
all_metrics = []
for pred, label in zip(preds, labels):
pred = (pred > threshold).astype(int)
accuracy = accuracy_score(label, pred)
precision = precision_score(label, pred, average='macro')
recall = recall_score(label, pred, average='macro')
f1_macro = f1_score(label, pred, average='macro')
f1_micro = f1_score(label, pred, average='micro')
all_metrics.append([accuracy, precision, recall, f1_macro, f1_micro])
avg_metrics = np.mean(all_metrics, axis=0)
print(avg_metrics)
return avg_metrics
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
for embedding_type in ['mdlm', 'esm', 'mlm']:
# Initialize datasets
train_dataset = LocalizationDataset(embedding_type,
hyperparams['train_data'],
hyperparams['esm_model_path'],
hyperparams['mlm_model_path'],
hyperparams['mdlm_model_path'],
device)
test_dataset = LocalizationDataset(embedding_type,
hyperparams['test_data'],
hyperparams['esm_model_path'],
hyperparams['mlm_model_path'],
hyperparams['mdlm_model_path'],
device)
# Prepare dataloaders
train_dataloader = DataLoader(train_dataset, batch_size=hyperparams["batch_size"], shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=hyperparams["batch_size"], shuffle=False)
# Initialize model, optimizer, and loss function
input_dim=640 if embedding_type=="mdlm" else 1280
model = LocalizationPredictor(input_dim=input_dim).to(device)
optimizer = optim.Adam(model.parameters(), lr=hyperparams["learning_rate"])
criterion = nn.BCELoss()
# Initialize main directory model checkpoints
base_checkpoint_dir = f"{path}/benchmarks/Supervised/Localization/model_checkpoints/{embedding_type}"
# Initialize subdirectory and name it based on hyperparameters
hyperparam_str = f"batch_{hyperparams['batch_size']}_lr_{hyperparams['learning_rate']}_epochs_{hyperparams['num_epochs']}"
model_checkpoint_dir = os.path.join(base_checkpoint_dir, hyperparam_str)
os.makedirs(model_checkpoint_dir, exist_ok=True)
# Training loop
for epoch in range(hyperparams["num_epochs"]):
# Train the model
train_loss = train(model, train_dataloader, optimizer, criterion, device)
print(f"EPOCH {epoch+1}/{hyperparams['num_epochs']}")
print(f"TRAIN LOSS: {train_loss:.4f}")
print("\n")
# Save the model checkpoint for the current epoch
checkpoint_path = os.path.join(model_checkpoint_dir, f"epoch{epoch + 1}.pth")
torch.save({
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': train_loss,
}, checkpoint_path)
print(f"Checkpoint saved at {checkpoint_path}\n")
# Save hyperparameters only once
if epoch == 0: # Hyperparameters don't change midway through training
hyperparams_file = os.path.join(model_checkpoint_dir, "hyperparams.txt")
with open(hyperparams_file, 'w') as f:
for key, value in hyperparams.items():
f.write(f"{key}: {value}\n")
print(f"Hyperparameters saved at {hyperparams_file}\n")
# Evaluate model on test dataset
print("Test set")
test_preds, test_labels = evaluate(model, test_dataloader, device)
test_metrics = calculate_metrics(test_preds, test_labels)
print(test_metrics)
print("TEST METRICS:")
print(f"Accuracy: {test_metrics[0]:.4f}")
print(f"Precision: {test_metrics[1]:.4f}")
print(f"Recall: {test_metrics[2]:.4f}")
print(f"F1 Macro Score: {test_metrics[3]:.4f}")
print(f"F1 Micro Score: {test_metrics[4]:.4f}")
#Save test results
test_results_file = os.path.join(model_checkpoint_dir, "test_results.txt")
with open(test_results_file, 'w') as f:
f.write("TEST METRICS:\n")
f.write(f"Accuracy: {test_metrics[0]:.4f}\n")
f.write(f"Precision: {test_metrics[1]:.4f}\n")
f.write(f"Recall: {test_metrics[2]:.4f}\n")
f.write(f"F1 Macro Score: {test_metrics[3]:.4f}\n")
f.write(f"F1 Micro: {test_metrics[4]:.4f}\n")
print(f"Test results saved at {test_results_file}\n") |