I am using pytorch for evaluating gradients of feed-forward network, but only for a subset of parameters, related to the first two layers. Since backpropagation is carried backwards layer by layer, I wonder: why is it computationally faster than evaluating gradients of whole network?
1 Answers
Pytorch builds a computation graph for backward propagation that only contains the minimum nodes and edges to get the accumulated gradient for leaves that require gradient. Even if the first two layers require gradient, there are many tensors (intermediate tensors or frozen parameters tensors) that are unused and that are cut in the backward graph. Plus the built-in function AccumulatedGradient that stores the gradients in .grad attribute is call less time reducing the total computation time too.
Here you can see an example for an "AddBackward Node" where for instance A is an intermediate tensor computed with the first two layers and B is the 3rd (constant) layer that can be ignored.
An other example: if you have a matrix-matrix product (MmBackward Node) that uses an intermediate tensor that not depends on the 2 first layers. In this case the tensor itself is required to compute the backprop but the "previous" tensors that were used to compute it can be ignored in the graph.
To visualize the sub-graph that is actually computed (and compare when the model is unfrozen), you can use torchviz.
Related