PyTorch Datasets

PyTorch Datasets#

The satrain.pytorch.datasets module provides PyTorch-compatible dataset classes for both tabular and spatial data.

Dataset Classes#

Subset Size Considerations#

When working with PyTorch datasets, subset size affects memory usage and training time:

Memory-efficient subsets:

  • xs, s - Can typically fit in memory, good for development

  • m - May require data chunking depending on available RAM

  • l, xl - Require careful memory management and data streaming

Recommended workflow:

  1. Prototype with xs or s

  2. Validate with m

  3. Train production models with l or xl

Usage Examples#

Tabular dataset for pixel-based models:

from satrain.pytorch.datasets import SatRainTabular
from torch.utils.data import DataLoader

# Start with small subset for development
dataset = SatRainTabular(
    base_sensor='gmi',
    split='training',
    subset='s',  # Small subset for development
    inputs=['gmi'],
    transforms=None
)

# Create DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

Spatial dataset for image-based models:

from satrain.pytorch.datasets import SatRainSpatial

# Use medium subset for spatial models
dataset = SatRainSpatial(
    base_sensor='gmi',
    split='training',
    subset='m',  # Medium subset for spatial work
    inputs=['gmi', 'geo'],
    tile_size=128,
    transforms=None
)

Progressive scaling workflow:

# 1. Start development with extra small
dev_dataset = SatRainTabular(subset='xs')

# 2. Validate with small to medium
val_dataset = SatRainTabular(subset='s')

# 3. Train final model with large
train_dataset = SatRainTabular(subset='l')

Memory management for large subsets:

from torch.utils.data import DataLoader

# For large subsets, use smaller batch sizes
if subset in ['l', 'xl']:
    batch_size = 16  # Smaller batches for large datasets
    num_workers = 2  # Fewer workers to reduce memory usage
else:
    batch_size = 32  # Standard batch size for smaller subsets
    num_workers = 4

dataloader = DataLoader(
    dataset,
    batch_size=batch_size,
    num_workers=num_workers,
    pin_memory=True if torch.cuda.is_available() else False
)

Complete Module Reference#