Add tests

Author: sgreenbury

tests/decoders/test_channels_last.py

@@ -0,0 +1,14 @@
+import torch
+
+from auto_cast.decoders.channels_last import ChannelsLast
+
+
+def test_channels_last_reorders_dimensions():
+    decoder = ChannelsLast()
+    x = torch.randn(2, 3, 4, 5, 6)
+
+    output = decoder(x)
+
+    assert output.shape == (2, 4, 5, 6, 3)
+    # Spot-check a value to ensure permutation matches expectation
+    assert torch.allclose(output[0, 0, 0, 0, 0], x[0, 0, 0, 0, 0])

tests/encoders/test_permute_concat.py

@@ -0,0 +1,55 @@
+import torch
+
+from auto_cast.encoders.permute_concat import PermuteConcat
+from auto_cast.types import Batch
+
+
+def _make_batch(
+    batch_size: int = 1,
+    t: int = 1,
+    w: int = 2,
+    h: int = 3,
+    c: int = 2,
+    const_c: int = 1,
+    scalar_c: int = 1,
+) -> Batch:
+    input_fields = torch.arange(batch_size * t * w * h * c, dtype=torch.float32)
+    input_fields = input_fields.view(batch_size, t, w, h, c)
+    output_fields = torch.zeros(batch_size, t, w, h, c)
+    constant_fields = torch.ones(batch_size, w, h, const_c)
+    constant_scalars = torch.full((batch_size, scalar_c), 5.0)
+    return Batch(
+        input_fields=input_fields,
+        output_fields=output_fields,
+        constant_scalars=constant_scalars,
+        constant_fields=constant_fields,
+    )
+
+
+def test_permute_concat_with_constants():
+    encoder = PermuteConcat(with_constants=True)
+    batch = _make_batch()
+
+    encoded = encoder(batch)
+
+    expected = batch.input_fields.permute(0, 4, 1, 2, 3)
+
+    base_channels = batch.input_fields.shape[-1]
+    assert batch.constant_fields is not None
+    assert batch.constant_scalars is not None
+    const_channels = batch.constant_fields.shape[-1]
+    scalar_channels = batch.constant_scalars.shape[-1]
+
+    assert encoded.shape == (
+        batch.input_fields.shape[0],
+        base_channels + const_channels + scalar_channels,
+        batch.input_fields.shape[1],
+        batch.input_fields.shape[2],
+        batch.input_fields.shape[3],
+    )
+
+    assert torch.allclose(encoded[:, :base_channels, ...], expected)
+    const_slice = encoded[:, base_channels : base_channels + const_channels, ...]
+    assert torch.allclose(const_slice, torch.ones_like(const_slice))
+    scalar_slice = encoded[:, -scalar_channels:, ...]
+    assert torch.allclose(scalar_slice, torch.full_like(scalar_slice, 5.0))

tests/models/test_encoder_processor_decoder.py

@@ -0,0 +1,80 @@
+import lightning as L
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+
+from auto_cast.decoders.channels_last import ChannelsLast
+from auto_cast.encoders.permute_concat import PermuteConcat
+from auto_cast.models.encoder_decoder import EncoderDecoder
+from auto_cast.models.encoder_processor_decoder import EncoderProcessorDecoder
+from auto_cast.nn.base import Module
+from auto_cast.types import Batch, Tensor
+
+
+class TinyProcessor(Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.conv = nn.Conv3d(in_channels=1, out_channels=1, kernel_size=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(x)
+
+
+def _toy_batch(
+    batch_size: int = 2,
+    t_in: int = 2,
+    t_out: int | None = None,
+    w: int = 4,
+    h: int = 4,
+    c: int = 1,
+) -> Batch:
+    t_out = t_out or t_in
+    input_fields = torch.randn(batch_size, t_in, w, h, c)
+    output_fields = torch.randn(batch_size, t_out, w, h, c)
+    return Batch(
+        input_fields=input_fields,
+        output_fields=output_fields,
+        constant_scalars=None,
+        constant_fields=None,
+    )
+
+
+class _BatchDataset(Dataset):
+    def __len__(self) -> int:
+        return 2  # keep small
+
+    def __getitem__(self, idx: int) -> Batch:
+        return _toy_batch(batch_size=1)
+
+
+dummy_loader = DataLoader(
+    _BatchDataset(),
+    batch_size=1,
+    collate_fn=lambda items: items[0],
+    num_workers=0,
+)
+
+
+def test_encoder_processor_decoder_training_step_runs():
+    encoder = PermuteConcat(with_constants=False)
+    decoder = ChannelsLast()
+    loss = nn.MSELoss()
+    encoder_decoder = EncoderDecoder(encoder=encoder, decoder=decoder, loss_func=loss)
+
+    processor = TinyProcessor()
+    model = EncoderProcessorDecoder.from_encoder_processor_decoder(
+        encoder_decoder=encoder_decoder,
+        processor=processor,
+        loss_func=loss,
+    )
+
+    batch = _toy_batch()
+    train_loss = model.training_step(batch, 0)
+
+    assert train_loss.shape == ()
+    train_loss.backward()
+
+    trainer = L.Trainer(
+        max_epochs=1, logger=False, enable_checkpointing=False, limit_train_batches=1
+    )
+    trainer.fit(model, train_dataloaders=dummy_loader, val_dataloaders=dummy_loader)