Why doesn't the AutoEncoder tutorial use Lightning?

[taimoorsohail]

Hi everyone,

I am building an ocean-dataset-based Autoencoder off the Autoencoder tutorial - very helpful thank you!

But I am curious why the tutorial as it is uses PET for setting up the pipeline, but then defines the Autoencoder and trains using torch.nn and related modules. Based on my read of current PET implementation, we could instead define a Lightning wrapper like (ignore the masked implementation, that is specific to my oceans routine):

import lightning as L
import torch
import torch.nn as nn
import torch.optim as optim

class LitAutoEncoder(L.LightningModule):
    def __init__(self, model, mask, lr=1e-4):
        super().__init__()
        self.model = model
        self.lr = lr
        self.criterion = nn.L1Loss()

        # Persist mask with device movement/checkpoints
        self.register_buffer(
            "mask",
            torch.as_tensor(mask, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        )

    def forward(self, x):
        # x: [B, C, H, W]
        mask = self.mask.expand(x.shape[0], 1, x.shape[2], x.shape[3])
        return self.model(x, mask)

    def training_step(self, batch, batch_idx):
        x = x[:, 0]
        if isinstance(x, (tuple, list)):
            x = x[0]

        x_hat = self.forward(x)

        # optional: ocean-only loss
        mask = self.mask.expand_as(x[:, :1])
        loss = torch.abs(x_hat - x)
        loss = (loss * mask).sum() / mask.sum().clamp_min(1.0)

        self.log("train_loss", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x = batch

        if isinstance(x, (tuple, list)):
            x = x[0]

        x_hat = self.forward(x)

        mask = self.mask.expand_as(x[:, :1])
        loss = torch.abs(x_hat - x)
        loss = (loss * mask).sum() / mask.sum().clamp_min(1.0)

        self.log("valid_loss", loss, prog_bar=True)
        return loss

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

then define a base model and data module as:

base_model = AutoEncoder(
    input_channel_count=2,
    output_channel_count=2,
)

lightning_model = LitAutoEncoder(
    model=base_model,
    mask=mask.values,
    lr=1e-4,
)

datamodule = pyearthtools.training.data.lightning.PipelineLightningDataModule(
    pipeline_i,
    **splits,
    batch_size=8,
    num_workers=0,
)

and then trains and fits:

trainer = pyearthtools.training.lightning.Train(
    lightning_model,
    datamodule,
    path="/g/data/v46/txs156/OM2-emulator/data/",
    trainer_kwargs={
        "max_epochs": 10,
        "num_sanity_val_steps": 1,
    },
)

trainer.fit()

Am I missing something here? Perhaps the aim is to keep the tutorial focussed only on pipeline preparation and not the other implemented PET routines?

[stevehadd]

I think there are two reasons why, fro my perspective (@ten Lee may wish to comment):

• For the momentum partnership hackathon, I wanted a more explicit training loop to show how the pipeline was being used, and the helper functionality of training hid away some of that
• It also reflects my lack of familiarity with lightning compared to basic PyTorch. Having similarly been revisiting the example in recent weeks I've come to the conclusion we should rather replace it with the lighting training function to make the most of pet to reduce boilerplate, so would be great if you created a or to update the autoencoder tutorial to use the lightning module stuff

[tennlee]

We don't really have set best practises for the model classes yet. The Lightning wrappers are just one way to go that seemed good at the time, but not all models necessarily fit. So from my perspective it's more about supporting variety. In terms of those particular tutorials, I didn't really want to slow down the hackathon by introducing too many concepts, as Stephen said and as you suggested. I also wanted to make it clear the PET pipelines are actually standard iterators which can be used with any model framework and there's no framework lock-in involved. However, I do think we need to expand out our docs for good practises for writing model classes.

[taimoorsohail]

@tennlee @stevehadd thanks for clarifying.

I do have an implementation that adapts the Autoencoder to use Lightning here: https://github.com/PaulSpence/OM2-emulator/blob/ts/train-on-more-than-one-timestep/notebooks/AutoEncoder_om2.ipynb

Let me know if you would like me to give back to the docs/elsewhere this implementation. You will notice though that eager YAML saving has substantially slowed down the current PET implementation, so this might need to be patched/modified if we want to implement a Lightning integration (see issue #265 ).