AI GitHub
PyTorch for tensor programs and neural-network training
Build and train neural networks with tensors, automatic differentiation and modules.
Key Takeaways
- PyTorch's tensor and autograd APIs support imperative training loops.
- Install a build that matches the device and start with one small tensor program.
- Train a small classifier and verify one gradient and parameter update.
- Precision, device and random seeds influence reproducibility.
How it works
PyTorch's tensor and autograd APIs support imperative training loops. Modules, optimizers and data utilities add structure around the model and its updates.
What you need
Install a build that matches the device and start with one small tensor program. Record dtype, device and input shape when checking numerical differences.
What to try first
Train a small classifier and verify one gradient and parameter update. Save and restore the model and compare predictions on the same batch.
Compare evaluation before and after saving while the model is in evaluation mode. Training-only behavior such as dropout can obscure a good checkpoint round trip.
Where to look for failures
Precision, device and random seeds influence reproducibility. Verify the supported installation for your accelerator and avoid loading serialized objects from untrusted sources.
Related Franklin pages
Keep the task and input representation fixed as you move between a model library and tensor frameworks. Separate compile or initialization time from subsequent execution.
- Compose transformations for differentiation, vectorization and compilation around array functions. JAX for transforming numerical programs
- Use Triton kernels for supported training operations and integrations. Liger Kernel for studying efficient language-model training operations
For the role of multilayer networks and how they learn, see Franklin’s glossary. Deep learning
Project documentation
Try it for your workflow
Explore the project’s documentation for setup and usage, or save this page to a collection for later.
Read the project documentation ↗