AI summary

This paper introduces an automatic differentiation module of PyTorch, a library designed to enable rapid research on machine learning models. It focuses on differentiation of purely imperative programs, with a focus on extensibility and low overhead. It achieves high performance on both CPU and GPU.

Abstract

In this article, we describe an automatic differentiation module of PyTorch a library designed to enable rapid research on machine learning models. It builds upon a few projects, most notably Lua Torch, Chainer, and HIPS Autograd [4], and provides a high performance environment with easy access to automatic differentiation of models executed on different devices (CPU and GPU). To make prototyping easier, PyTorch does not follow the symbolic approach used in many other deep learning frameworks, but focuses on differentiation of purely imperative programs, with a focus on extensibility and low overhead. Note that this preprint is a draft of certain sections from an upcoming paper covering all PyTorch features.