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

Three Python deep learning frameworks, three strengths: PyTorch for research and new architectures, Keras for a high-level API, JAX for compiled code on TPUs.

How to choose:

• Research and new model architectures: PyTorch
• High-level model building on JAX or PyTorch: Keras
• NumPy-style array code compiled for GPUs and TPUs: JAX
• PyTorch training without the loop boilerplate: PyTorch Lightning
• Reinforcement learning environments: Gymnasium
• Reinforcement learning algorithms: Stable-Baselines3

PyTorch models are nn.Module subclasses, and autograd builds the computational graph as your code runs. Install it with the command from its selector, which matches your OS and GPU. To keep a model, save its state_dict instead of pickling the whole module. Wrap your model in torch.compile to speed it up with minimal code changes. To train on more than one GPU, use DistributedDataParallel.

Keras is a multi-framework API: a Keras model can run as a PyTorch Module or as a JAX function. Install a backend next to it and set KERAS_BACKEND before you import Keras. Build simple models as a Sequential stack of layers and anything more complex with the functional API. Save the whole model to a .keras file with model.save() instead of pickling it. The file reloads with any backend.

JAX does accelerator-oriented array computation with a NumPy-style API and composable transformations: jax.grad for derivatives, jax.jit for compilation, and jax.vmap for batching. The transformations only work on functionally pure functions, so pass all data in as arguments and return every result. JAX itself stays narrow: to train neural networks, use the JAX AI Stack, with Flax NNX for models and Optax for optimizers. On NVIDIA GPUs, the JAX team strongly recommends installing CUDA and cuDNN from pip wheels.

PyTorch Lightning organizes PyTorch code to remove boilerplate: you write the model logic in a LightningModule, and the Trainer handles devices, precision, and distributed training. Install it as the lightning package. Its style guide recommends keeping each LightningModule self-contained, the model separate from the system that trains it, and data loading in a LightningDataModule. To keep your own training loop, use Lightning Fabric, which scales a plain PyTorch script after you change a few lines.

Gymnasium is an API standard for reinforcement learning, with a collection of reference environments. It's the maintained fork of OpenAI's Gym, and many older tutorials still use Gym's old API, so follow its migration guide when you port one. Register your own environment so gymnasium.make() creates it like a built-in one, and run check_env on it to catch common issues.

Stable-Baselines3 is a set of reliable implementations of reinforcement learning algorithms in PyTorch, and it trains on any environment that follows the Gymnasium interface. It assumes you know some reinforcement learning. Its tips page recommends starting from the RL Zoo's tuned hyperparameters and normalizing the agent's input. Evaluate the agent on a separate test environment, since training adds exploration noise. Pick an algorithm by your action space first: DQN handles only discrete actions, and SAC only continuous ones.

PyTorch's security policy says running untrusted models is equivalent to running untrusted code, so run untrusted ones in a sandbox. Load checkpoints with weights_only=True in torch.load, and leave safe_mode on when Keras loads a model.