Defining learning models compatible with PyRelationAL

To interact with the PyRelationAL library, models need to be wrapped within a PyRelationAL model manager (pyrelational.model_managers) that defines required methods for instantiation, training, and testing.

Using a pytorch module

Let’s first look an an example model manager for a simple pytorch Module

import torch
import torch.nn as nn
from pyrelational.model_managers.mcdropout_model import MCDropoutModelManager


class TorchModuleWrapper(MCDropoutModelManager):
    def __init__(self,
        model_class,
        model_config,
        trainer_config,
        n_estimators= 5,
        eval_dropout_prob= 0.2,
    ):
        super(Wrapper, self).__init__(model_class,model_config,trainer_config,n_estimators,eval_dropout_prob)

    def train(self, train_loader, val_loader=None):
        model = self._init_model()
        criterion = nn.MSELoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=self.trainer_config["lr"])
        for _ in range(self.trainer_config["epochs"]):
            for x, y in train_loader:
                optimizer.zero_grad()
                out = model(x)
                loss = criterion(out, y)
                loss.backward()
                optimizer.step()
        self._current_model = model # store current model

    def test(self, loader):
        if not self.is_trained():
            raise ValueError("No current model, call 'train(train_loader, valid_loader)' to train the model first")
        criterion = nn.MSELoss()
        self._current_model.eval()
        with torch.no_grad():
            tst_loss, cnt = 0, 0
            for x, y in loader:
                loss = criterion(self._current_model(x), y)
                cnt += x.size(0)
                tst_loss += loss.item()
        return {"test_loss": tst_loss/cnt}

This can now be instantiated with any pytorch module, for instance

model = TorchModuleWrapper(
                  nn.Linear,
                  {"in_features": 5, "out_features":1},
                  {"epochs":100, "lr":3e-4},
                  n_estimators=10,
               )

Using pytorch lightning module

PyRelationAL implements default classes (see pyrelational.model_managers.lightning_model.LightningModelManager()) relying on pytorch lightning as the Trainer class offload much of the training routine definition to pytorch lighntning. For example, users can create ensembles of pytorch lightning modules directly as

import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score
from lightning.pytorch import LightningModule
from pyrelational.model_managers.ensemble_model_manager import LightningEnsembleModelManager

# step 1: define the LightningModule with necessary methods
class DigitClassifier(LightningModule):
    """Custom module for a simple convnet classifier"""

    def __init__(self, dropout_rate=0, lr=3e-4):
        super(DigitClassifier, self).__init__()
        self.layer_1 = nn.Linear(8*8, 16)
        self.layer_2 = nn.Linear(16, 32)
        self.dropout = nn.Dropout(dropout_rate)
        self.layer_3 = nn.Linear(32, 10)
        self.lr = lr

    def forward(self, x):
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.layer_3(x)
        x = F.log_softmax(x, dim=1)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        self.log("loss", loss.item())
        return loss

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        self.log("test_loss", loss)

        # compute accuracy
        _, y_pred = torch.max(logits.data, 1)
        accuracy = accuracy_score(y, y_pred)
        self.log("accuracy", accuracy)

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

# step 2: wrap the model in a LightningEnsembleModelManager
wrapper = LightningEnsembleModelManager(
              DigitClassifier,
              {"dropout_rate":0.1, "lr":3e-4},
              {"epochs":1,"accelerator":"gpu", "devices":1},
              n_estimators=5,
        )

See the examples folder in the source repository for more examples.