Use preprocessed and cached tfrecord datasets with PyTorch

[ampersandmcd]

• IceNet version: 0.2.8
• Python version: 3.11.7
• Operating System: CentOS Linux 7 (BAS HPC)

Description

To accelerate training with PyTorch, it would be useful to make use of the preprocessed and cached tfrecord datasets we've produced for operational Tensorflow training! Loading from these tfrecord files is (should be?) much faster than loading from NetCDF files on the fly using the generate_sample logic in, e.g., the DaskMultiWorkerLoader.generate_sample defined in data/loaders/dask.py which must perform a decent amount of computation.

What I Did

I've implemented an IterableIceNetDatasetPyTorch class which inherits from torch.utils.data.IterableDataset as a first go, and will link it below in a pull request. The following script demonstrates its use and can easily be stepped through in a debugger. The batching logic works properly with torch.utils.data.DataLoader in the script below, but I've seen some weird behaviour (overrunning the number of samples that should be generated per epoch) during training runs when num_workers > 1 in the torch.utils.data.DataLoader so there's room to improve this first implementation.

import os
import torch
from utils import IterableIceNetDataSetPyTorch

dataset_config = "dataset_config.exp23_south.json"

ds = IterableIceNetDataSetPyTorch(dataset_config, "test", batch_size=4, shuffling=False)
dl = torch.utils.data.DataLoader(ds, batch_size=4, shuffle=False)

for i, batch in enumerate(dl):
    x, y, sw = batch
    print(x, y, sw)
    print(i)

[JimCircadian]

CC @bnubald who has done some work in this area already