Saved tensors#

Some operations need intermediary results to be saved during the forward passin order to execute the backward pass. For example, the function$x\mapsto x^2$x↦x2 saves the input$x$x to compute the gradient.

When defining a custom PythonFunction, you can usesave_for_backward() to savetensors during the forward pass andsaved_tensors to retrieve themduring the backward pass. SeeExtending PyTorch for more information.

For operations that PyTorch defines (e.g.torch.pow()), tensors areautomatically saved as needed. You can explore (for educational or debuggingpurposes) which tensors are saved by a certaingrad_fn by looking for itsattributes starting with the prefix_saved.

x=torch.randn(5,requires_grad=True)y=x.pow(2)print(x.equal(y.grad_fn._saved_self))# Trueprint(xisy.grad_fn._saved_self)# True

In the previous code,y.grad_fn._saved_self refers to the same Tensor object asx.But that may not always be the case. For instance:

x=torch.randn(5,requires_grad=True)y=x.exp()print(y.equal(y.grad_fn._saved_result))# Trueprint(yisy.grad_fn._saved_result)# False

Under the hood, to prevent reference cycles, PyTorch haspacked the tensorupon saving andunpacked it into a different tensor for reading. Here, thetensor you get from accessingy.grad_fn._saved_result is a different tensorobject thany (but they still share the same storage).

Whether a tensor will be packed into a different tensor object depends onwhether it is an output of its owngrad_fn, which is an implementation detailsubject to change and that users should not rely on.

You can control how PyTorch does packing / unpacking withHooks for saved tensors.

Division by Zero in Autograd#

When performing division by zero in PyTorch (e.g.,x/0), the forward pass will produceinf values following IEEE-754 floating point arithmetic. While theseinf values can be masked out before computing the final loss (e.g., via indexing or masking), the autograd system still tracks and differentiates through the full computation graph, including the division by zero operation.

During backpropagation, this can lead to problematic gradient expressions. For example:

x=torch.tensor([1.,1.],requires_grad=True)div=torch.tensor([0.,1.])y=x/div# Results in [inf, 1]mask=div!=0# [False, True]loss=y[mask].sum()loss.backward()print(x.grad)# [nan, 1], not [0, 1]

In this example, even though we only use the masked output (which excludes the division by zero), autograd still computes gradients through the full computation graph, including the division by zero operation. This results innan gradients for the masked elements, which can cause training instability.

To avoid this issue, there are several recommended approaches:

1. Mask before division:

x=torch.tensor([1.,1.],requires_grad=True)div=torch.tensor([0.,1.])mask=div!=0safe=torch.zeros_like(x)safe[mask]=x[mask]/div[mask]loss=safe.sum()loss.backward()# Produces safe gradients [0, 1]

2. Use MaskedTensor (experimental API):

fromtorch.maskedimportas_masked_tensorx=torch.tensor([1.,1.],requires_grad=True)div=torch.tensor([0.,1.])y=x/divmask=div!=0loss=as_masked_tensor(y,mask).sum()loss.backward()# Cleanly handles "undefined" vs "zero" gradients

The key principle is to prevent the division by zero operation from being recorded in the computation graph, rather than masking its results after the fact. This ensures that autograd only computes gradients through valid operations.

This behavior is important to keep in mind when working with operations that might produceinf ornan values, as masking the outputs does not prevent the problematic gradients from being computed.

Settingrequires_grad#

requires_grad is a flag, defaulting to falseunless wrappedin ann.Parameter, that allows for fine-grained exclusion ofsubgraphs from gradient computation. It takes effect in both theforward and backward passes:

During the forward pass, an operation is only recorded in the backward graph ifat least one of its input tensors require grad.During the backward pass (.backward()), only leaf tensors withrequires_grad=True will have gradients accumulated into their.gradfields.

It is important to note that even though every tensor has this flag,setting it only makes sense for leaf tensors (tensors that do not have agrad_fn, e.g., ann.Module’s parameters).Non-leaf tensors (tensors that do havegrad_fn) are tensors that have abackward graph associated with them. Thus their gradients will be neededas an intermediary result to compute the gradient for a leaf tensor thatrequires grad. From this definition, it is clear that all non-leaf tensorswill automatically haverequire_grad=True.

Settingrequires_grad should be the main way you control which partsof the model are part of the gradient computation, for example, if you need tofreeze parts of your pretrained model during model fine-tuning.

To freeze parts of your model, simply apply.requires_grad_(False) tothe parameters that you don’t want updated. And as described above,since computations that use these parameters as inputs would not be recorded inthe forward pass, they won’t have their.grad fields updated in the backwardpass because they won’t be part of the backward graph in the first place, asdesired.

Because this is such a common pattern,requires_grad can also be set atthe module level withnn.Module.requires_grad_().When applied to a module,.requires_grad_() takes effect on allof the module’s parameters (which haverequires_grad=True by default).

Grad Modes#

Apart from settingrequires_grad there are also three grad modes that canbe selected from Python that can affect how computations in PyTorch areprocessed by autograd internally: default mode (grad mode), no-grad mode,and inference mode, all of which can be togglable via context managers anddecorators.

Default Mode (Grad Mode)#

The “default mode” is the mode we are implicitly in when no other modes likeno-grad and inference mode are enabled. To be contrasted with“no-grad mode” the default mode is also sometimes called “grad mode”.

The most important thing to know about the default mode is that it is the onlymode in whichrequires_grad takes effect.requires_grad is always overriddento beFalse in both the two other modes.

No-grad Mode#

Computations in no-grad mode behave as if none of the inputs require grad.In other words, computations in no-grad mode are never recorded in the backward grapheven if there are inputs that haverequire_grad=True.

Enable no-grad mode when you need to perform operations that should not berecorded by autograd, but you’d still like to use the outputs of thesecomputations in grad mode later. This context manager makes it convenient todisable gradients for a block of code or function withouthaving to temporarily set tensors to haverequires_grad=False, and thenback toTrue.

For example, no-grad mode might be useful when writing an optimizer: whenperforming the training update you’d like to update parametersin-place without the update being recorded by autograd.You also intend to use the updated parameters for computations ingrad mode in the next forward pass.

The implementations intorch.nn.init alsorely on no-grad mode when initializing the parameters as to avoidautograd tracking when updating the initialized parameters in-place.

Inference Mode#

Inference mode is the extreme version of no-grad mode. Just like in no-gradmode, computations in inference mode are not recorded in the backward graph, butenabling inference mode will allow PyTorch to speed up your model even more.This better runtime comes with a drawback: tensors created in inference modewill not be able to be used in computations to be recorded by autograd afterexiting inference mode.

Enable inference mode when you are performing computations that do not haveinteractions with autograd, AND you don’t plan on using the tensors createdin inference mode in any computation that is to be recorded by autograd later.

It is recommended that you try out inference mode in the parts of your codethat do not require autograd tracking (e.g., data processing and model evaluation).If it works out of the boxfor your use case it’s a free performance win. If you run into errors afterenabling inference mode, check that you are not using tensors created ininference mode in computations that are recorded by autograd after exiting inferencemode. If you cannot avoid such use in your case, you can always switch backto no-grad mode.

For details on inference mode please seeInference Mode.

For implementation details of inference mode seeRFC-0011-InferenceMode.

Evaluation Mode (nn.Module.eval())#

Evaluation mode is not a mechanism to locally disable gradient computation.It is included here anyway because it is sometimes confused to be such a mechanism.

Functionally,module.eval() (or equivalentlymodule.train(False)) are completelyorthogonal to no-grad mode and inference mode. Howmodel.eval() affectsyour model depends entirely on the specific modules used in your model andwhether they define any training-mode specific behavior.

You are responsible for callingmodel.eval() andmodel.train() if yourmodel relies on modules such astorch.nn.Dropout andtorch.nn.BatchNorm2d that may behavedifferently depending on training mode, for example, to avoid updating yourBatchNorm running statistics on validation data.

It is recommended that you always usemodel.train() whentraining andmodel.eval() when evaluating your model (validation/testing) evenif you aren’t sure your model has training-mode specific behavior, because amodule you are using might be updated to behave differently in training andeval modes.

In-place correctness checks#

Every tensor keeps a version counter, that is incremented every time it ismarked dirty in any operation. When a Function saves any tensors for backward,a version counter of their containing Tensor is saved as well. Once you accessself.saved_tensors it is checked, and if it is greater than the saved valuean error is raised. This ensures that if you’re using in-placefunctions and not seeing any errors, you can be sure that the computedgradients are correct.

Concurrency on CPU#

When you runbackward() orgrad() via python or C++ API in multiplethreads on CPU, you are expecting to see extra concurrency instead ofserializing all the backward calls in a specific order during execution(behavior before PyTorch 1.6).

Non-determinism#

If you are callingbackward() from multiple threads concurrently and haveshared inputs (i.e. Hogwild CPU training), then non-determinism should be expected.This can occur because parameters are automatically shared across threads,as such, multiple threads may access and try to accumulate the same.gradattribute during gradient accumulation. This is technically not safe, andit might result in race condition and the result might be invalid to use.

Users developing multithreaded models featuring shared parameters should have thethreading model in mind and should understand the issues described above.

The functional APItorch.autograd.grad() may be used to calculate thegradients instead ofbackward() to avoid non-determinism.

Graph retaining#

If part of the autograd graph is shared between threads, i.e. run firstpart of forward single thread, then run second part in multiple threads,then the first part of graph is shared. In this case different threadsexecutegrad() orbackward() on the same graph might have issue ofdestroying the graph on the fly of one thread, and the other thread willcrash in this case. Autograd will error out to the user similar to what callbackward() twice with outretain_graph=True, and let the user knowthey should useretain_graph=True.

Thread Safety on Autograd Node#

Since Autograd allows the caller thread to drive its backward execution forpotential parallelism, it’s important that we ensure thread safety on CPU withparallelbackward() calls that share part/whole of the GraphTask.

Custom Pythonautograd.Functions are automatically thread safe because of GIL.For built-in C++ Autograd Nodes (e.g. AccumulateGrad, CopySlices) and customautograd::Functions, the Autograd Engine uses thread mutex locking to ensurethread safety on autograd Nodes that might have state write/read.

No thread safety on C++ hooks#

Autograd relies on the user to write thread safe C++ hooks. If you want the hookto be correctly applied in multithreading environment, you will need to writeproper thread locking code to ensure the hooks are thread safe.

What are complex derivatives?#

The mathematical definition of complex-differentiability takes thelimit definition of a derivative and generalizes it to operate oncomplex numbers. Consider a function$f:\mathrm{C}\to \mathrm{C}$f:C→C,

$$f(z=x+yj)=u(x,y)+v(x,y)j$$f(z=x+yj)=u(x,y)+v(x,y)j

where$u$u and$v$v are two variable real valued functionsand$j$j is the imaginary unit.

Using the derivative definition, we can write:

$$f^{\mathrm{\prime }}(z)=\underset{h\to 0,h\in C}{\lim }\frac{f(z+h)-f(z)}{h}$$f′(z)=h→0,h∈Clim​hf(z+h)−f(z)​

In order for this limit to exist, not only must$u$u and$v$v must bereal differentiable, but$f$f must also satisfy the Cauchy-Riemannequations. Inother words: the limit computed for real and imaginary steps ($h$h)must be equal. This is a more restrictive condition.

The complex differentiable functions are commonly known as holomorphicfunctions. They are well behaved, have all the nice properties thatyou’ve seen from real differentiable functions, but are practically of nouse in the optimization world. For optimization problems, only real valued objectivefunctions are used in the research community since complex numbers are not part of anyordered field and so having complex valued loss does not make much sense.

It also turns out that no interesting real-valued objective fulfill theCauchy-Riemann equations. So the theory with holomorphic function cannot beused for optimization and most people therefore use the Wirtinger calculus.

Wirtinger Calculus comes into the picture …#

So, we have this great theory of complex differentiability andholomorphic functions, and we can’t use any of it at all, because manyof the commonly used functions are not holomorphic. What’s a poormathematician to do? Well, Wirtinger observed that even if$f(z)$f(z)isn’t holomorphic, one could rewrite it as a two variable function$f(z,z\ast )$f(z,z∗) which is always holomorphic. This is because real andimaginary of the components of$z$z can be expressed in terms of$z$z and$z^{\ast }$z∗ as:

Re(z)Im(z)​=2z+z∗​=2jz−z∗​​

Wirtinger calculus suggests to study$f(z,z^{\ast })$f(z,z∗) instead, which isguaranteed to be holomorphic if$f$f was real differentiable (anotherway to think of it is as a change of coordinate system, from$f(x,y)$f(x,y)to$f(z,z^{\ast })$f(z,z∗).) This function has partial derivatives$\frac{\mathrm{\partial }}{\mathrm{\partial }z}$∂z∂​ and$\frac{\mathrm{\partial }}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂​.We can use the chain rule to establish arelationship between these partial derivatives and the partialderivatives w.r.t., the real and imaginary components of$z$z.

∂x∂​∂y∂​​=∂x∂z​∗∂z∂​+∂x∂z∗​∗∂z∗∂​=∂z∂​+∂z∗∂​=∂y∂z​∗∂z∂​+∂y∂z∗​∗∂z∗∂​=1j∗(∂z∂​−∂z∗∂​)​

From the above equations, we get:

∂z∂​∂z∗∂​​=1/2∗(∂x∂​−1j∗∂y∂​)=1/2∗(∂x∂​+1j∗∂y∂​)​

which is the classic definition of Wirtinger calculus that you would find onWikipedia.

There are a lot of beautiful consequences of this change.

• For one, the Cauchy-Riemann equations translate into simply saying that$\frac{\mathrm{\partial }f}{\mathrm{\partial }{z}^{\ast }}=0$∂z∗∂f​=0 (that is to say, the function$f$f can be writtenentirely in terms of$z$z, without making reference to$z^{\ast }$z∗).

• Another important (and somewhat counterintuitive) result, as we’ll see later, is that when we do optimization on a real-valued loss, the step we shouldtake while making variable update is given by$\frac{\mathrm{\partial }Loss}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂Loss​ (not$\frac{\mathrm{\partial }Loss}{\mathrm{\partial }z}$∂z∂Loss​).

For more reading, check out:https://arxiv.org/pdf/0906.4835.pdf

How is Wirtinger Calculus useful in optimization?#

Researchers in audio and other fields, more commonly, use gradientdescent to optimize real valued loss functions with complex variables.Typically, these people treat the real and imaginary values as separatechannels that can be updated. For a step size$\alpha \mathrm{/}2$α/2 and loss$L$L, we can write the following equations in${\mathrm{R}}^2$R2:

xn+1​yn+1​​=xn​−(α/2)∗∂x∂L​=yn​−(α/2)∗∂y∂L​​

How do these equations translate into complex space$\mathrm{C}$C?

zn+1​​=xn​−(α/2)∗∂x∂L​+1j∗(yn​−(α/2)∗∂y∂L​)=zn​−α∗1/2∗(∂x∂L​+j∂y∂L​)=zn​−α∗∂z∗∂L​​

Something very interesting has happened: Wirtinger calculus tells usthat we can simplify the complex variable update formula above to onlyrefer to the conjugate Wirtinger derivative$\frac{\mathrm{\partial }L}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂L​, giving us exactly the step we take in optimization.

Because the conjugate Wirtinger derivative gives us exactly the correct step for a real valued loss function, PyTorch gives you this derivativewhen you differentiate a function with a real valued loss.

How does PyTorch compute the conjugate Wirtinger derivative?#

Typically, our derivative formulas take ingrad_output as an input,representing the incoming Vector-Jacobian product that we’ve alreadycomputed, aka,$\frac{\mathrm{\partial }L}{\mathrm{\partial }{s}^{\ast }}$∂s∗∂L​, where$L$Lis the loss of the entire computation (producing a real loss) and$s$s is the output of our function. The goal here is to compute$\frac{\mathrm{\partial }L}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂L​, where$z$z is the input ofthe function. It turns out that in the case of real loss, we canget away withonly calculating$\frac{\mathrm{\partial }L}{\mathrm{\partial }{s}^{\ast }}$∂s∗∂L​,even though the chain rule implies that we also need tohave access to$\frac{\mathrm{\partial }L}{\mathrm{\partial }s}$∂s∂L​. If you wantto skip this derivation, look at the last equation in this sectionand then skip to the next section.

Let’s continue working with$f:\mathrm{C}\to \mathrm{C}$f:C→C defined as$f(z)=f(x+yj)=u(x,y)+v(x,y)j$f(z)=f(x+yj)=u(x,y)+v(x,y)j. As discussed above,autograd’s gradient convention is centered around optimization for realvalued loss functions, so let’s assume$f$f is a part of largerreal valued loss function$g$g. Using chain rule, we can write:

(1)#$$\frac{\mathrm{\partial }L}{\mathrm{\partial }{z}^{\ast }}=\frac{\mathrm{\partial }L}{\mathrm{\partial }u}\ast \frac{\mathrm{\partial }u}{\mathrm{\partial }{z}^{\ast }}+\frac{\mathrm{\partial }L}{\mathrm{\partial }v}\ast \frac{\mathrm{\partial }v}{\mathrm{\partial }{z}^{\ast }}$$∂z∗∂L​=∂u∂L​∗∂z∗∂u​+∂v∂L​∗∂z∗∂v​

Now using Wirtinger derivative definition, we can write:

$$\begin{array}{r}{\displaystyle \frac{\mathrm{\partial }L}{\mathrm{\partial }s}=1\mathrm{/}2\ast (\frac{\mathrm{\partial }L}{\mathrm{\partial }u}-\frac{\mathrm{\partial }L}{\mathrm{\partial }v}j)}\\ {\displaystyle \frac{\mathrm{\partial }L}{\mathrm{\partial }{s}^{\ast }}=1\mathrm{/}2\ast (\frac{\mathrm{\partial }L}{\mathrm{\partial }u}+\frac{\mathrm{\partial }L}{\mathrm{\partial }v}j)}\end{array}$$∂s∂L​=1/2∗(∂u∂L​−∂v∂L​j)∂s∗∂L​=1/2∗(∂u∂L​+∂v∂L​j)​

It should be noted here that since$u$u and$v$v are realfunctions, and$L$L is real by our assumption that$f$f is apart of a real valued function, we have:

(2)#$${\left(\frac{\mathrm{\partial }L}{\mathrm{\partial }s}\right)}^{\ast }=\frac{\mathrm{\partial }L}{\mathrm{\partial }{s}^{\ast }}$$(∂s∂L​)∗=∂s∗∂L​

i.e.,$\frac{\mathrm{\partial }L}{\mathrm{\partial }s}$∂s∂L​ equals to$grad\mathrm{\_}outpu{t}^{\ast }$grad_output∗.

Solving the above equations for$\frac{\mathrm{\partial }L}{\mathrm{\partial }u}$∂u∂L​ and$\frac{\mathrm{\partial }L}{\mathrm{\partial }v}$∂v∂L​, we get:

(3)#$$\begin{array}{r}{\displaystyle \frac{\mathrm{\partial }L}{\mathrm{\partial }u}=\frac{\mathrm{\partial }L}{\mathrm{\partial }s}+\frac{\mathrm{\partial }L}{\mathrm{\partial }{s}^{\ast }}}\\ {\displaystyle \frac{\mathrm{\partial }L}{\mathrm{\partial }v}=1j\ast (\frac{\mathrm{\partial }L}{\mathrm{\partial }s}-\frac{\mathrm{\partial }L}{\mathrm{\partial }{s}^{\ast }})}\end{array}$$∂u∂L​=∂s∂L​+∂s∗∂L​∂v∂L​=1j∗(∂s∂L​−∂s∗∂L​)​

Substituting(3) in(1), we get:

∂z∗∂L​​=(∂s∂L​+∂s∗∂L​)∗∂z∗∂u​+1j∗(∂s∂L​−∂s∗∂L​)∗∂z∗∂v​=∂s∂L​∗(∂z∗∂u​+∂z∗∂v​j)+∂s∗∂L​∗(∂z∗∂u​−∂z∗∂v​j)=∂s∂L​∗∂z∗∂(u+vj)​+∂s∗∂L​∗∂z∗∂(u+vj)∗​=∂s∂L​∗∂z∗∂s​+∂s∗∂L​∗∂z∗∂s∗​​

Using(2), we get:

(4)#∂z∗∂L​​=(∂s∗∂L​)∗∗∂z∗∂s​+∂s∗∂L​∗(∂z∂s​)∗=(grad_output)∗∗∂z∗∂s​+grad_output∗(∂z∂s​)∗​​

This last equation is the important one for writing your own gradients,as it decomposes our derivative formula into a simpler one that is easyto compute by hand.

How can I write my own derivative formula for a complex function?#

The above boxed equation gives us the general formula for allderivatives on complex functions. However, we still need tocompute$\frac{\mathrm{\partial }s}{\mathrm{\partial }z}$∂z∂s​ and$\frac{\mathrm{\partial }s}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂s​.There are two ways you could do this:

• The first way is to just use the definition of Wirtinger derivatives directly and calculate$\frac{\mathrm{\partial }s}{\mathrm{\partial }z}$∂z∂s​ and$\frac{\mathrm{\partial }s}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂s​ byusing$\frac{\mathrm{\partial }s}{\mathrm{\partial }x}$∂x∂s​ and$\frac{\mathrm{\partial }s}{\mathrm{\partial }y}$∂y∂s​(which you can compute in the normal way).

• The second way is to use the change of variables trick and rewrite$f(z)$f(z) as a two variable function$f(z,z^{\ast })$f(z,z∗), and computethe conjugate Wirtinger derivatives by treating$z$z and$z^{\ast }$z∗ as independent variables. This is often easier; for example, if the function in question is holomorphic, only$z$z will be used (and$\frac{\mathrm{\partial }s}{\mathrm{\partial }{z}^{\ast }}$∂z∗∂s​ will be zero).

Let’s consider the function$f(z=x+yj)=c\ast z=c\ast (x+yj)$f(z=x+yj)=c∗z=c∗(x+yj) as an example, where$c\in \mathrm{R}$c∈R.

Using the first way to compute the Wirtinger derivatives, we have.

Using(4), andgrad_output = 1.0 (which is the default grad output value used whenbackward() is called on a scalar output in PyTorch), we get:

$$\frac{\mathrm{\partial }L}{\mathrm{\partial }{z}^{\ast }}=1\ast 0+1\ast c=c$$∂z∗∂L​=1∗0+1∗c=c

Using the second way to compute Wirtinger derivatives, we directly get:

∂z∂s​∂z∗∂s​​=∂z∂(c∗z)​=c=∂z∗∂(c∗z)​=0​