atomai
atomai copied to clipboard
Custom dataloaders for AtomAI models
Passing a custom data loader in AtomAI models can be a useful feature. This would allow using e.g. Kornia data augmentation pipelines. It could look like this
segmodel = aoi.models.Segmentor(nb_classes=3)
segmodel.fit(train_loader=custom_train_loader, test_loader=custom_test_loader). # instead of passing X_train, y_train, X_test, y_test
It will also be helpful to have a utility function that builds a dataloader with Kornia data augmentation functions.
I will add the custom data loader functionality, starting with the Segmentor class.