Designing the training pipeline

[leifdenby]

@arjj8 and I had some good discussions yesterday about how to design the training pipeline API. We want to make sure that the API is flexible enough to support different use cases, while also being easy to use and understand.

The first thing we did was try to go through the steps of data transformation, from the source datasets (assuming the datastructure of that mlcast-dataset-validator imposes) to the final torch.Tensor that is fed into the model.

From the LDCast arhitecture (work by @franchg, @arjj8, @mfroelund, @martinbo-meteo, links below), we have the following steps:

1. read source dataset
2. tile this dataset in space and time (e.g. 12 time steps, 128x128 spatial tiles), eliminating tiles with too much missing data
3. sample tiles to normalize the distribution of training samples (e.g. ensuring that we have a good distribution of samples across different precipitation intensities)
4. normalize the data for training (e.g. scaling precipitation values to a certain range)
5. batch the data into torch.Tensor batches for training

We have summarised this in the following flowchart:

flowchart TB
    A["mlcast-datasets intake repo<br />zarr URL"]
    H["user-provided zarr dataset URL"]
    B["source dataset<br />shape: [n_time, x, y]"]
    C["Tiled xr.DataArray<br />shape: [tile_id, n_time_window, x_tile, y_tile]"]
    D["Sampled xr.DataArray<br />shape: [sampled_tile_id, n_time_sample, x_tile, y_tile]"]
    I["Normalized xr.DataArray<br />shape: [sampled_tile_id, n_time_sample, x_tile, y_tile]"]
    E["torch.Tensor<br />shape: [batch_size, n_time_sample, x_tile, y_tile]"]
    F["CSV index<br />{tiling_id}.samples.csv"]
    G["CSV index<br />{sampling_id}.samples.csv"]

    A -- "open dataset:<br />OpenMLCastDataset()" --> B
    H -- "open dataset:<br />OpenXarrayDataset()" --> B
    B -- "tiling step:<br />TilingSampler()" --> C
    C -- "write tiling index:<br />{tiling_id}.samples.csv<br />tiling_id = f\"{dataset_id}.{tile_size}.{n_time_window}\"" --> F
    C -- "sampling step:<br />BinNormSampler()" --> D
    D -- "write samples index:<br />{sampling_id}.samples.csv<br />sampling_id = f\"{tiling_id}.{n_time_sample}\"" --> G
    D -- "normalize step:<br />NormalizeForTraining()" --> I
    I -- "batching step" --> E

Loading

Because some of these steps might be optional (maybe some architectures might not require tiling) or we might them done in a different way (different sampling strategy for example) I think it would be nice to have a design where we can define which steps to apply and in which order.

# Example for LDCast architecture
pipeline = [
    MLCastCatalogDataset("dmi.precipitation.5min", var_name="rainrate"),
    TilingSampler(
        n_time_window=12,
        tile_size=(128, 128),
    ),
    BinNormSampler(
        aggregation="sum",
        n_scalar_total_bins=10,
    ),
    NormalizeForTraining(),
    ToTorchTensor()
]

# Example without sampling step
pipeline = [
    MLCastCatalogDataset("dmi.precipitation.5min", var_name="rainrate"),
    TilingSampler(
        n_time_window=12,
        tile_size=(128, 128),
    ),
    NormalizeForTraining(),
    ToTorchTensor()
]

I had hoped to use torchdata datapipes for this, because torch data pipes allow for exactly this kind of composability for constructing data processing pipelines. However I have since yesterday found out that datapipes were depricated from pytorch in July 2024 and the last version to include them is torchdata==0.9.0 (which is supported by torch==2.5.0, but not torch==2.6.0, so we'd have to stick to torch==2.5.x if we wanted to use data pipes), and although a replacement is being worked on, it isn't clear to me where things are at.

One question here is: How much of this computation do we want ahead of training time? Currently, I think @franchg you are taking the view that we should:

1. use the source data directly (i.e. not create a training dataset)
2. create the tiling index and sampling indexes ahead of training
3. use fixed scaling for normalization - can we get away with this in general?

Does this seem about right @franchg? Also, what should we call this step? It seems like "data-preparation" to me, can we call it that?

Doing part of this processing, but not all, ahead of training makes the design a bit tricky here, but not impossible :)

What is the problem? At its core it comes down to mixing different implementations of parallelism:

1. When computing for example the tile statistics over the whole dataset for the sampling step, it would be natural to use dask with xarray to compute these statistics in parallel across the whole dataset.
2. When we want to use the pipeline during training, we want to be able to use torch.utils.data.DataLoader to load batches of data in parallel using multiple workers.

As far as I can tell, there isn't a good way to mix these two types of parallelism in a single pipeline (I think this is why torch datapipes were removed, and there is no obvious alternative yet). If we for example had multiple workers doing data-loading with torch.utils.data.DataLoader, we can't then produce the tile statistics for the sampling step using dask in parallel across the whole dataset, because each worker would only see a subset of the data.

So where does this leave us? I have asked @arjj8 to refactor @franchg's tiling and sampling code as pure xarray code that can be run in parallel with dask, which return the resulting tile-set as xr.DataArray and has the side effect of creating the tiling and sampling CSV files. With this we can then either 1) have separate calls to create the processing pipeline during data-preparation (i.e. either creating a tiled+sampled dataset on disk, or using the indexes during training to load the relevant tiles and samples from the source dataset) and training, 2) still use torch datapipes and restrict ourselves to torch<2.6.0, I'm not sure if this would be an issue, or 3) or come up with an abstraction where we can use dask during `pre-rain

All this discussion might all be overkill at this stage. But I think it is worth considering how we design this, because once we want to experiment with different approaches and combine different datasets this could all become quite complex, and we want to make sure that we have a design that is flexible enough to support different use cases, while also being easy to use and understand.

Links to relevant discussions and code:

• Adding the LDCast architecture #4
• LDCast implementation in MLCast (in progress) #8
• current implementation of tiling and sampling: https://github.com/mlcast-community/mlcast-dataset-sampler

Thoughts?

[arjj8]

hej!

So, another idea that I was thinking of is using fiddle package which let you build multitude functions, (tilling,sampled, etc..) and keep them as objects until the calculations are needed. Additionally, it facilitates the modification of parameters inside the objects. I would need to check how this work with parallel, but I think there is potential.

For now I can start with the daskand xarrayapproach and see how well works.

[mfroelund]

One idea came to my mind, which I'm not sure really makes sense to follow, but I still want to share it in case you can see anything useful in it.

So I wondered if it would make sense to make a kind of running estimate of the tile statistics, which is updated for every batch. So let's say we load the first batch and compute the tile statistics of that batch. Then this tile statistics would of coarse be completely off compared to the tile statistics of the complete dataset. But if we save the information about what tiles we sample in the first batch, and in the second batch update the tile statistics from the first batch with that of the second batch, we could then choose tiles in the second batch such that we evened out a bit the distribution of tiles from the first batch. As training progress, the tile statistics would gradually approach that of the complete dataset.

If this makes sense at all, one immediate question that I wonder about, is how this influence the training, i.e. how does it influence training that the distribution of tiles in the beginning of the training is off compared to if you took the tile statistics of the whole dataset into account from the beginning?
Also, how this would work in parallel with some kind of shared tile statistics information, I don't know.

[martinbo-meteo]

Hi all,

What do you mean exactly by 'tile statistics' ? Is it what is done by the tiling and sampling code of @franchg ?

I might be missing the whole point here 😄 but I do not really see the problem to pass first through the full dataset as is done for the moment by @franchg's code: if we plan to train the models for more than 100 epochs, is it a problem to load in memory the dataset once more ? (if the tiled samples only represent a fraction of the full training dataset, this preprocessing will take a time equivalent to that of a few epochs). We could still save the indices in csv files as is done for the moment if we want to resume the training. Again, I am maybe missing something...

About the normalization, I was thinking that it might be interesting to use a fixed scaling: I think it makes more sense to do it this way if we want to run a model on a dataset on which it has not been trained (it seems meaningless to me to feed a model with dataset A scaled with max(A) if it has been trained on dataset B scaled with max(B)). We could even use a scaling common to all datasets and models ?

PS: I think that the samplers mentioned in LDCast (#8) have nothing to do with the sampling we are discussing here :)