SR-Forge

Collation and the Batch Convention¶

What is Collation?¶

Collation is the process of combining multiple individual samples into a single batched structure that can be processed efficiently by a neural network. SR-Forge uses stack-based collation: each sample stores values in their natural shape (e.g. a single image is [C, H, W]), and collation adds a new batch dimension via torch.stack:

e1 = Entry(image=torch.randn(3, 64, 64))  # [C, H, W] — natural shape
e2 = Entry(image=torch.randn(3, 64, 64))

batch = Entry.collate([e1, e2])
# batch.image.shape == [2, 3, 64, 64]  — batch dim added by stacking

An Entry is a flexible container that holds tensors, lists, dicts, and nested structures. Collation handles all of these types, turning a list of Entries into one batched Entry where every field is combined in a type-appropriate way — tensors are stacked, lists are collected, and nested structures are recursed into.

The Stack-Based Convention¶

SR-Forge follows a convention where single samples use natural shapes — no batch dimension. A single image is [C, H, W], a name is "scene_042" (a bare string), a band list is ["b2", "b8"]. Collation adds the batch dimension — see the full rules table below for every type.

When a Dataset returns an Entry, fields should be in their natural form:

class MyDataset(Dataset):
    def __getitem__(self, idx) -> Entry:
        image = load_image(path)           # Returns [C, H, W] — natural shape
        target = load_image(target_path)
        return Entry(
            name=self.names[idx],          # Bare string
            image=image,                   # [C, H, W]
            target=target,
            bands=["b2", "b8"],            # Flat list
            scale=4,                       # Bare int
        )

The `is_batched` Flag¶

Collation marks the result with _is_batched = True. This flag controls how transforms dispatch:

Batched entry (is_batched == True): Data has a batch dimension. Tensors are [B, C, H, W].
Unbatched entry (is_batched == False): Data is in natural form. Tensors are [C, H, W].

You don't need to manage this flag manually — collate() sets it, and integer indexing (batch[i]) clears it.

How Collation Works¶

Collation combines a list of Entries into a single batched Entry. It happens automatically inside DataLoader via SR-Forge's built-in GeneralCollation collate function:

from srforge.data.loader import DataLoaderFactory

loader = DataLoaderFactory(dataset, batch_size=4).get_loader()

for batch in loader:
    # batch is a single Entry with stacked tensors
    print(batch.image.shape)  # [4, 3, 64, 64]

The key operation is stacking: each sample has shape [C, H, W], so stacking four of them produces [4, C, H, W]:

Sample 1:  [3, 64, 64]  --+
Sample 2:  [3, 64, 64]    +--  torch.stack(dim=0)  -->  [4, 3, 64, 64]
Sample 3:  [3, 64, 64]    |
Sample 4:  [3, 64, 64]  --+

Entry.collate() also stores the batch size and sets is_batched = True on the result. See Batch Indexing for how to extract individual samples back out.

Collation Rules by Type¶

Field type	Single sample	After collation (B=4)
Tensor (same shape)	`[C, H, W]`	`[4, C, H, W]` (stacked)
Tensor (different shapes)	varies	kept as list (not stackable)
str	`"scene_042"`	`["scene_042", "scene_043", ...]` (collected)
int / float	`4`	`tensor([4, 4, ...])`
bool	`True`	`tensor([True, False, ...], dtype=torch.bool)`
list	`["b2", "b8"]`	`[["b2", "b8"], ["b2", "b8"], ...]` (collected)
tuple	`(t1, t2)`	`[(t1, t2), (t3, t4)]` (collected)
dict	`{"a": t1, "b": t2}`	`{"a": stack(t1s), "b": stack(t2s)}` (recurse)
None (all)	`None`	`[None, None, ...]` (collected)
None (mixed with values)	`None` or value	kept as list of raw values

Tensors¶

Tensors of the same shape are stacked along a new dim 0. If shapes differ across the batch (e.g. variable-size images), collation falls back to returning a plain list — you'll need to handle padding or grouping yourself.

Lists and Tuples¶

Lists and tuples are collected — each sample's value becomes an element in the batch list:

e1 = Entry(bands=["b2", "b8"])
e2 = Entry(bands=["b2", "b8"])

batch = Entry.collate([e1, e2])
# batch.bands == [["b2", "b8"], ["b2", "b8"]]

For multi-temporal data where each sample has a sequence of images:

e1 = Entry(lrs=[t1, t2])      # 1 sample, 2 timesteps
e2 = Entry(lrs=[t3, t4, t5])  # 1 sample, 3 timesteps

batch = Entry.collate([e1, e2])
# batch.lrs == [[t1, t2], [t3, t4, t5]]  — a list of 2 inner lists

Dicts¶

Dict fields are collated recursively — each key is collated independently:

e1 = Entry(bands={"rgb": torch.randn(3, 64, 64), "nir": torch.randn(2, 64, 64)})
e2 = Entry(bands={"rgb": torch.randn(3, 64, 64), "nir": torch.randn(2, 64, 64)})

batch = Entry.collate([e1, e2])
# batch.bands["rgb"].shape == [2, 3, 64, 64]
# batch.bands["nir"].shape == [2, 2, 64, 64]

None¶

None fields are always collected into a list — both when all entries are None and when some are mixed with actual values. This preserves batch size information and keeps indexing consistent (list[i] works the same way as for strings or other non-stackable types):

e1 = Entry(mask=None)
e2 = Entry(mask=torch.randn(1, 64, 64))

batch = Entry.collate([e1, e2])
# batch.mask == [None, tensor([1, 64, 64])]  — a list of raw values

Warning

All non-None values in a field must be the same type. Mixing different types (e.g. a string in one sample and an int in another) raises a TypeError — only None gets special treatment as "missing data".

Bare Scalars¶

Bare strings, ints, floats, and bools are handled automatically:

# Strings are collected into a list
e1 = Entry(name="scene_042")
e2 = Entry(name="scene_043")
batch = Entry.collate([e1, e2])
# batch.name == ["scene_042", "scene_043"]

# Ints and floats become tensors
e1 = Entry(scale=2)
e2 = Entry(scale=4)
batch = Entry.collate([e1, e2])
# batch.scale == tensor([2, 4])

# Bools become bool tensors
e1 = Entry(flag=True)
e2 = Entry(flag=False)
batch = Entry.collate([e1, e2])
# batch.flag == tensor([True, False])

Writing Transforms for Batched Data¶

Batched transforms (transform()) should assume tensors have a batch dimension. Use ... (ellipsis) indexing when you want to operate on spatial dimensions regardless of how many leading dimensions exist:

class CropBorder(DataTransform):
    def __init__(self, border: int):
        super().__init__()
        self.border = border

    def transform(self, image: torch.Tensor) -> torch.Tensor:
        b = self.border
        return image[..., b:-b, b:-b]  # Works for any number of leading dims

For operations that need explicit batch/channel dimensions, unpack them directly:

class PerChannelNormalize(DataTransform):
    def transform(self, image: torch.Tensor) -> torch.Tensor:
        B, C, H, W = image.shape
        mean = image.mean(dim=(2, 3), keepdim=True)
        std = image.std(dim=(2, 3), keepdim=True)
        return (image - mean) / (std + 1e-8)

Transforms can also be written to work on single samples — implement transform_unbatched() instead, and the framework handles batching and unbatching for you (more on this in Transforms):

class PerSampleResize(DataTransform):
    def transform_unbatched(self, image: torch.Tensor) -> torch.Tensor:
        # image is [C, H, W] — no batch dim
        return F.interpolate(image.unsqueeze(0), size=self.size).squeeze(0)

Reversing Collation: Batch Indexing¶

Collation is reversible. Given a batched Entry, you can extract individual samples using integer indexing:

batch = Entry.collate([e1, e2, e3])
sample = batch[0]         # Unbatched Entry — batch dim removed
samples = batch.unbatch() # [batch[0], batch[1], batch[2]]

Integer indexing uses true indexing (tensor[i]) — the batch dimension is removed (PyTorch semantics). Slice indexing (batch[0:1]) preserves the batch dimension and the is_batched flag:

# Integer index: removes batch dim
sample = batch[0]
# sample.image.shape == [3, 64, 64]  — no batch dim
# sample.name == "scene_042"         — bare string
# sample.is_batched == False

# Slice: preserves batch dim
sub = batch[0:1]
# sub.image.shape == [1, 3, 64, 64]  — batch dim preserved
# sub.name == ["scene_042"]           — list
# sub.is_batched == True

Round-Trip Invariant¶

# collate → index recovers original values
e1 = Entry(name="s1", image=tensor_1)
e2 = Entry(name="s2", image=tensor_2)
batch = Entry.collate([e1, e2])

r1 = batch[0]
# r1.name == "s1"
# torch.equal(r1.image, e1.image)

For the full API and round-trip invariants, see Entry: Batch Indexing.

GraphEntry¶

GraphEntry uses PyTorch Geometric's own collation via torch_geometric.loader.DataLoader. This means graph data is batched by concatenating node features and adjusting edge indices — the standard PyG approach. SR-Forge's DataLoaderFactory automatically selects the right loader based on the Entry type.

Next: Datasets — How to load data and create custom datasets

Collation and the Batch Convention¶

What is Collation?¶

The Stack-Based Convention¶

The is_batched Flag¶

How Collation Works¶

Collation Rules by Type¶

Tensors¶

Lists and Tuples¶

Dicts¶

None¶

Bare Scalars¶

Writing Transforms for Batched Data¶

Reversing Collation: Batch Indexing¶

Round-Trip Invariant¶

GraphEntry¶

The `is_batched` Flag¶