Composition of operations in PyTorch

Consider the start of the forward method from the PyTorch tutorial:

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
		...
		return x

In the operation described, x is a torch.Tensor object. The various methods invoked return torch.Tensor objects. Hence, the network can be described as a chained sequence of tensor-to-tensor mappings. So we can un-nest the forward operation as the following:

def forward(self, x: torch.Tensor) -> torch.Tensor:
	x = self.conv1(x)
	x = F.relu(x)
	x = F.max_pool2d(x, (2, 2))
	...
	return x

In English, we’d say “To make a forward pass, we first apply our first of two learned 2-dimensional convolution operation to the input tensor, then we apply a rectified linear operation to the output. We downsample this result using 2x2 max pooling…”