sample.py

import os

import lightning as L

import torch

from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

from jigsaw import piece
from jigsaw import composite

PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
BATCH_SIZE = 256 if torch.cuda.is_available() else 64


FEATURE_MNIST = "mnist"
OUTPUT_CLASSIFICATION = "classification"
LABEL_CLASSIFICATION = "label"


class Linear(piece.Piece):
    def __init__(self):
        super().__init__(piece_type="module", name="linear")
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def inputs(self) -> tuple[str, ...]:
        return tuple([FEATURE_MNIST])

    def outputs(self) -> tuple[str, ...]:
        return tuple([OUTPUT_CLASSIFICATION])

    def forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        mnist = inputs[FEATURE_MNIST]
        out = torch.relu(self.l1(mnist))
        output = {OUTPUT_CLASSIFICATION: out}
        return output


class MNISTModel(L.LightningModule):
    def __init__(self, model: Linear, loss: piece.Piece):
        super().__init__()
        self.model = model
        self.loss = loss
        self.composite = composite.Composite([model, loss])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        inputs = {FEATURE_MNIST: x.view(x.size(0), -1)}
        outputs = self.model.forward(inputs)
        return outputs[OUTPUT_CLASSIFICATION]

    def training_step(
        self, batch: tuple[torch.Tensor, torch.Tensor], batch_nb: int
    ) -> torch.Tensor:
        x, y = batch
        inputs = {FEATURE_MNIST: x.view(x.size(0), -1), LABEL_CLASSIFICATION: y}
        outputs = self.composite.forward(inputs)
        return outputs[self.loss.name]

    def configure_optimizers(self) -> torch.optim.Optimizer:
        return torch.optim.Adam(self.parameters(), lr=0.02)


if __name__ == "__main__":
    # Init our model
    linear = Linear()
    loss = piece.Loss(
        loss_fn=torch.nn.CrossEntropyLoss(),
        loss_input_names=[OUTPUT_CLASSIFICATION, LABEL_CLASSIFICATION],
    )
    mnist_model = MNISTModel(linear, loss)

    # Init DataLoader from MNIST Dataset
    train_ds = MNIST(
        PATH_DATASETS, train=True, download=True, transform=transforms.ToTensor()
    )
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, num_workers=7)

    # Initialize a trainer
    trainer = L.Trainer(
        accelerator="auto",
        devices=1,
        max_epochs=3,
    )

    # Train the model ⚡
    trainer.fit(mnist_model, train_loader)