modules.py

from typing import List
import torch.nn as nn
from torch.optim import Adam
import pytorch_lightning as pl

class Route2Vec(pl.LightningModule):
    def __init__(self, encoder: nn.Module, classifiers: nn.ModuleList, criteria: List, loss_weights: List, label_names: List, lr: float, batch_size: int) -> None:
        super().__init__()
        self.encoder = encoder
        self.classifiers = classifiers
        self.lr = lr
        self.batch_size = batch_size
        self.criteria = criteria
        self.loss_weights = loss_weights
        self.label_names = label_names
        self.save_hyperparameters(ignore=['encoder', 'classifiers'])

    def forward(self, x):
        _, h = self.encoder(x)
        outs = [cla(h) for cla in self.classifiers]
        return outs, h

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=self.lr)

    def training_step(self, batch, batch_idx):
        loss = self.forward_iteration(batch, batch_idx, "train")
        return {'loss': loss}

    def validation_step(self, batch, batch_idx):
        return self.forward_iteration(batch, batch_idx, "val")

    def forward_iteration(self, batch, batch_idx, mode):
        x = batch[0]
        labels = batch[1:]
        _, h = self.encoder(x)
        preds = [cla(h) for cla in self.classifiers]
        loss_list = [criterion(pred, label) for criterion, pred, label in zip(self.criteria, preds, labels)]

        for loss, label_name in zip(loss_list, self.label_names):
            self.log(f"{mode}_loss_{label_name}", loss, prog_bar=False, sync_dist=True)
        
        total_loss = sum([loss*weight for loss, weight in zip(loss_list, self.loss_weights)])
        self.log(f"{mode}_loss", total_loss, prog_bar=True, sync_dist=True)

        return total_loss

if __name__ == "__main__":
    from architecture import TransformerEncoder
    from architecture import DenseClassifier
    enc = TransformerEncoder(True, True, "none", 1.0, 36, 512, 8, 1)
    clas = nn.ModuleList([DenseClassifier(512,1), DenseClassifier(512,6)])
    criteria = [nn.MSELoss(), nn.CrossEntropyLoss()]
    loss_weights = [0.5, 0.5]
    model = Route2Vec(enc, clas, criteria, loss_weights, lr=1e-4, batch_size=128)
    print(model)
    params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"{params:,}")