Introduce a training API about axon HOT 4 CLOSED

FWIW I like this breakout - it combines the best of PyTorch's flexibility and provides logical callbacks during the training process. Do you see the eval_objective providing a mechanism for early stopping, or should that be another callback?

In line with what many other ML libs do, I agree that Axon should define a Protocol for datasets that any external data management library could implement support for. I also think including a basic in-memory implementation for testing, quick ramp-up, and providing a reference example would add great value.

IMO PyTorch has one of the most reasonable interfaces for working with custom datasets; I'm actually in the process of migrating one project from TF to PyTorch due to the complexity of implementing custom dataset generators.

from axon.

@arpieb Yes! I actually think it would be best to adjust after_epoch to take the evaluation state as well so early stopping can be implemented as a callback. That would also be useful for implementing model checkpointing based on validation metrics.

I've really only ever worked in TensorFlow/TFDS, but I have heard there is more of a preference for PyTorch for some of these things. I haven't really considered what kind of dataset integration Axon should have, mainly because I figured we would want to create a more general abstraction that works for the whole ecosystem, but it's definitely worth some thought to have something convenient.

from axon.

I think it's going to be hard to come up with a general abstraction for all use cases of Nx tensor collections. For example, Axon SL models will rely on sample-label pairs, which is not a requirement for most other data ops. A dataframe equivalent to pandas (exdas? nexdas? ;)) would need to be columnar to handle the mixture of numeric (Nx tensor) and non-numeric columns.

Ultimately all Python-based ML solutions seem to fall back to requiring a collection (len, get_item) or iterable (iter) to satisfy dataset needs, which in turn usually return one or more numpy arrays. Even DL4j forces any kind of transform pipeline or generator to ultimately spit up something implementing the INDArray interface (their attempt at replicating numpy.ndarray) before it can be handed off to a model.

Many UL models will not require labels, so 2+ rank tensors would fill the need just fine for them.

from axon.

This looks a lot like PyTorch Lightning! And that's a compliment. As much as I have distaste for much of the Python deep learning ecosystem, I think PTL has emerged from some early teething problems to be one of the most sane training APIs out there. Axon could do much worse than to emulate their approach to callbacks.

I agree wholeheartedly with @arpieb's take here (and in elixir-nx/nx#301) that defining some protocol for the inevitable dataframe libraries is essential here. Dataloading and its complexities have always been a bugbear for Pytorch and Tensorflow.

from axon.