Using TorchGeo within AzureML for model training

[PenguinJunk]

Hello, before we get mired down in solving the next steps in our work, we wonder if there is any experience in this community combining TorchGeo with training models in AzureML.

Setting the scene:

• We are using AzureML
• Our data are stored as "flightblocks", each flightblock represents a different geographical region and contains multiple images in a single directory, typically 100Gb in size
• We have created an AzureML Dataset for each flightblock (FileDataset with the v1 SDK)
• Our desire is to train on datasets comprising patches/chips sampled from the images in multiple flightblocks (such that the dataset represents multiple geographic regions)
• We would like this sampling of patches to occur on-the-fly, rather than pre-preparing and storing datasets of patches (because we want to experiment with patch parameters and so don't want to store lots of datasets)
• We have established that we can use TorchGeo to Union a small selection of data from more than one flightblock, and dataloaders to sample patches from across Unioned dataset

We are now beginning to train models in AzureML and wonder what experience others may have with training using TorchGeo dataloaders in AzureML, and if any 'best practice' has yet been established.

For instance:

• If we are distributing training, can we pass the TorchGeo dataloader into the training job so that AzureML handles batching etc.? (if so, how?)
• Alternatively, is the approach to pass a list of AzureML datasets OR image file paths into the training job and perform the Union and sampling on each distributed node?
• Has anyone worked on testing the distribution of sampled patches in the resulting batches?

We'll carry on plugging away at this but would be deeply grateful for any insights into how we best going about combing the power of TorchGeo data sampling with AzureML training.