Examples with public data - Moving MNIST

[jsta]

I am interested in testing some of these architectures with public data. Before I do that, does such an example already exist?

[jsta]

Potentially with Moving MNIST

[jsta]

Here's an attempt. I'm sure there's errors to be fixed and optimizations to be done, but it does run and loss goes down so that's something?

# python applications/harnesses/moving_mnist/train.py
# python applications/evaluation/TandVplot.py -S --basedir . -I 0 -Y 6000 -Nt 2 -Nv 250
import torch
import logging
import numpy as np
import torch.nn as nn
from tqdm import tqdm
from numpy.lib.stride_tricks import sliding_window_view
from torchvision.datasets import MovingMNIST
from torch.utils.data import Dataset

from yoke.models.vit.swin.bomberman import LodeRunner
import yoke.torch_training_utils as tr
import yoke.helpers.logger as yl


class mmnist_dataSet(Dataset):
    """Moving MNIST dataset."""

    def __init__(self, fraction=1, fraction_side="left") -> None:
        self.mmnist = MovingMNIST(".", download=True)
        total_len = 1000
        seq_len = 20
        pairs_per_seq = seq_len - 1
        frac_range = range(0, int(fraction * total_len))
        if fraction_side == "right":
            frac_range = range(int(fraction * total_len), total_len)
        self.seq_id = [
            x for xs in [np.repeat(i, pairs_per_seq) for i in frac_range] for x in xs
        ]
        pairs_local = [
            sliding_window_view(np.arange(0, seq_len), window_shape=2)
            for _ in frac_range
        ]
        self.pairs_local = np.concatenate(pairs_local)

    def __len__(self) -> int:
        return len(self.seq_id)

    def __getitem__(
        self, index: int
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Return a tuple of a batch's input and output data."""
        img_pair = self.mmnist[self.seq_id[index]][self.pairs_local[index]]
        start_img = (
            torch.tensor(np.expand_dims(img_pair[0, 0, ...], 0)).to(torch.float32) / 255
        )
        end_img = (
            torch.tensor(np.expand_dims(img_pair[1, 0, ...], 0)).to(torch.float32) / 255
        )
        Dt = torch.tensor(0.25, dtype=torch.float32)  # arbitrary value
        return start_img, end_img, Dt


if __name__ == "__main__":
    yl.configure_logger("yoke_logger", level=logging.INFO)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = LodeRunner(
        default_vars=["var1"],
        image_size=(64, 64),
        patch_size=(8, 8),
        embed_dim=4,
        emb_factor=2,
        num_heads=2,
        block_structure=(1, 1, 3, 1),
        window_sizes=[
            (4, 4),
            (4, 4),
            (2, 2),
            (1, 1),
        ],
        patch_merge_scales=[
            (2, 2),
            (2, 2),
            (2, 2),
        ],
    )

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=1e-3,
        betas=(0.9, 0.999),
        eps=1e-08,
        weight_decay=0.01,
    )

    # Use `reduction='none'` so loss on each sample in batch can be recorded.
    loss_fn = nn.MSELoss(reduction="none")

    model.to(device)

    # initialize outside of epoch loop because this is a single channel only
    train_dataset = mmnist_dataSet(0.75, "left")
    val_dataset = mmnist_dataSet(0.25, "right")

    train_dataloader = tr.make_dataloader(
        dataset=train_dataset,
        batch_size=2,
        num_batches=250,
        num_workers=1,
        prefetch_factor=2,
    )
    val_dataloader = tr.make_dataloader(
        dataset=val_dataset,
        batch_size=2,
        num_batches=25,
        num_workers=1,
        prefetch_factor=2,
    )

    num_epochs = 20
    for epochIDX in tqdm(range(num_epochs)):
        tr.train_simple_loderunner_epoch(
            [0],
            training_data=train_dataloader,
            validation_data=val_dataloader,
            model=model,
            optimizer=optimizer,
            loss_fn=loss_fn,
            epochIDX=epochIDX,
            train_per_val=10,
            train_rcrd_filename="train.csv",
            val_rcrd_filename="val.csv",
            device=device,
            verbose=False,
        )
        torch.cuda.empty_cache()

[hickmank]

@jsta Hey, I'm just getting around to actually looking at our "issues" tab and noticed you've been doing some great work.

Yes, please put in a harness example PR with moving MNIST. A summer student just put in a vanilla MNIST example but the more we have of those the better.

Happy to give advice if you submit a PR for this and it's not totally perfect yet. That's probably the easiest way to go back and forth about this code.