awesome-python/website/data/category_intros/deep-learning.md

PyTorch suits new models, since it debugs like plain Python. Gymnasium environments pass arrays to whichever Python deep learning framework trains your agent.

How to choose:

• Building and training neural networks: PyTorch
• RL environments, built-in or your own: Gymnasium
• A standard training loop scaled across GPUs: PyTorch Lightning
• NumPy-style code compiled for GPUs and TPUs: JAX
• One model API across JAX, TensorFlow, and PyTorch: Keras
• An existing TensorFlow codebase: TensorFlow
• Ready-made RL algorithms to train an agent: Stable Baselines3

PyTorch runs each line of code as it's reached, so a debugger and stack traces read the way you'd expect. Its quickstart shows the pieces of a project: a Dataset holds your samples and labels, and a DataLoader wraps it in an iterable. A model subclasses nn.Module, defines its layers in __init__, and passes data through them in forward. To keep a trained model, save its state_dict rather than pickling the whole model, which ties the file to your exact classes and directory layout.

PyTorch Lightning organizes PyTorch code so you write the model logic and it handles the engineering. A LightningModule is still an nn.Module, with your training step and configure_optimizers added, and the Trainer runs the loop around it. Moving to more GPUs, to TPUs, or to 16-bit precision is a Trainer argument, with no change to your model code.

JAX pairs a NumPy-like API with composable transformations: jax.grad takes gradients, jax.jit compiles, and jax.vmap vectorizes. They're designed for functionally pure functions, so pass every input as an argument and return every result. The core holds little that's specific to deep learning. For neural networks, add the JAX AI Stack: Flax for models, Optax for optimizers, and Orbax for checkpoints.

Gymnasium, a fork of OpenAI's Gym, is a standard API for reinforcement learning: one interface between learning algorithms and environments, plus reference environments that follow it. Its basic loop comes down to make(), Env.reset(), Env.step(), and Env.render(), and wrappers change an environment without touching its code. For your own problem, subclass gymnasium.Env with observation and action spaces, and register it so gymnasium.make() builds it like a built-in one.

Stable Baselines3 is a set of reliable implementations of RL algorithms in PyTorch, and it assumes you already know some RL. It follows a scikit-learn-like syntax: create a model such as A2C("MlpPolicy", env), call learn(), then predict(). It trains on any environment that follows the Gymnasium interface. Its RL tips pick the algorithm by action space, since DQN handles only discrete actions and SAC only continuous ones.

For more frameworks, papers, and courses, see awesome-deep-learning.