File size: 1,893 Bytes
6afa7ea e33b6c9 6afa7ea e33b6c9 6afa7ea e33b6c9 6afa7ea e33b6c9 6afa7ea | 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 | import pytorch_lightning as pl
import torch
import torch.nn as nn
from torch.optim import Adam
from model import LigandGNN, ProteinTransformer
class BindingAffinityModelPL(pl.LightningModule):
def __init__(self, num_node_features, hidden_channels_gnn, lr):
super().__init__()
self.save_hyperparameters() # Save hyperparameters for easy access
self.lr = lr
self.ligand_gnn = LigandGNN(
input_dim=num_node_features, hidden_channels=hidden_channels_gnn
)
self.protein_transformer = ProteinTransformer(vocab_size=26)
self.head = nn.Sequential(
nn.Linear(128 + 128, 256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, 1)
)
self.criterion = nn.MSELoss()
def forward(self, x, edge_index, batch, protein_seq):
ligand_vec = self.ligand_gnn(x, edge_index, batch)
batch_size = batch.max().item() + 1
protein_seq = protein_seq.view(batch_size, -1)
protein_vec = self.protein_transformer(protein_seq)
combined = torch.cat([ligand_vec, protein_vec], dim=1)
return self.head(combined)
def training_step(self, batch, batch_idx):
# We don't need .to(device), zero_grad, backward, PL handles that
out = self(batch.x, batch.edge_index, batch.batch, batch.protein_seq)
loss = self.criterion(out.squeeze(), batch.y.squeeze())
self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
return loss
def validation_step(self, batch, batch_idx):
out = self(batch.x, batch.edge_index, batch.batch, batch.protein_seq)
loss = self.criterion(out.squeeze(), batch.y.squeeze())
self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
return loss
def configure_optimizers(self):
return Adam(self.parameters(), lr=self.lr, weight_decay=1e-4)
|