A computation graph records the operations that transform inputs into outputs. In PyTorch, the graph can support two different jobs: calculating gradients during training and representing a model for export. section 5.5.1 section 17.2.2
A graph behind a loss
Set `requires_grad=True` on model parameters, compute a loss, and call `loss.backward()`. PyTorch follows the recorded operations backward and puts the loss derivatives in `params.grad`. If you repeat this in a training loop, zero the gradient first: backward calls accumulate gradients rather than replace them. section 5.5.1
For deployment, `torch.export` takes a `torch.nn.Module` and an example input and produces an ahead-of-time traced graph. Its nodes describe operations and the path data takes through the model. This exported representation serves a different purpose from the graph traversed by `backward()`. section 17.2.2 section 5.5.1
Listing 17.9 Examining exported program graph representation
# In[]
print(exported_program)
# Out[]
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, p_fc1_weight: "f32[5, 10]", p_fc1_bias: "f32[5]",
p_fc2_weight: "f32[2, 5]", p_fc2_bias: "f32[2]", x: "f32[2, 10]"):
# File: .py:10 in forward, code: x = self.fc1(x)
linear: "f32[2, 5]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);
x = p_fc1_weight = p_fc1_bias = None
# File: .py:11 in forward, code: x = self.relu(x)
relu: "f32[2, 5]" = torch.ops.aten.relu.default(linear); linear = None
# File: .py:12 in forward, code: x = self.fc2(x)
linear_1: "f32[2, 2]" = torch.ops.aten.linear.default(relu, p_fc2_weight, p_fc2_bias);
relu = p_fc2_weight = p_fc2_bias = None
return (linear_1,)
Inspect an exported program to see a concrete operation sequence: linear transformation, ReLU, then another linear transformation, with tensor shapes shown along the way. section 17.2.2
