The same network in PyTorch

Everything so far was NumPy, so that nothing was hidden. From here on we use PyTorch (Paszke et al. (2019)), the library the Lab's generated code is also written in. It gives us what NumPy does not: a way to differentiate a computation automatically. Chapters 3 and 5 depend on it. This section shows that the PyTorch network is the same object we built by hand.

Tensors

A PyTorch tensor is a NumPy array that can also remember how it was computed. The two convert freely:

import numpy as np
import torch

a = np.array([[1.0, 2.0], [3.0, 4.0]])
t = torch.from_numpy(a)                  # shares memory with the array
print(t)
print("dtype :", t.dtype)
print("matmul:", (t @ t).tolist())
print("a tensor made from scratch:", torch.ones(1).dtype)

tensor([[1., 2.],
        [3., 4.]], dtype=torch.float64)
dtype : torch.float64
matmul: [[7.0, 10.0], [15.0, 22.0]]
a tensor made from scratch: torch.float32

Note the dtype. NumPy's default is 64-bit floating point (float64); PyTorch's default for new tensors is 32-bit (float32). For ordinary machine learning 32 bits is plenty and twice as fast. For PINNs it is worth remembering that second derivatives of a network can lose several digits to rounding, so some PINN code switches to float64 with torch.set_default_dtype(torch.float64). We will say so when it matters.

A network as a module

The layers of Section 2.2 are torch.nn.Linear objects (each holds a weight of shape $(n_\text{out}, n_\text{in})$ and a bias, exactly our $W$ and $\mathbf{b}$), and the whole network is a torch.nn.Module whose forward method is our forward pass:

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, sizes):
        super().__init__()
        self.layers = nn.ModuleList(
            nn.Linear(n_in, n_out) for n_in, n_out in zip(sizes[:-1], sizes[1:])
        )

    def forward(self, x):
        for layer in self.layers[:-1]:
            x = torch.tanh(layer(x))          # hidden layers: bend
        return self.layers[-1](x)             # last layer: plain weighted sum

torch.manual_seed(0)
net = MLP([2, 4, 4, 1])
print(net)
print("parameters:", sum(p.numel() for p in net.parameters()))
for name, p in net.named_parameters():
    print(f"  {name:16s} {tuple(p.shape)}")

MLP(
  (layers): ModuleList(
    (0): Linear(in_features=2, out_features=4, bias=True)
    (1): Linear(in_features=4, out_features=4, bias=True)
    (2): Linear(in_features=4, out_features=1, bias=True)
  )
)
parameters: 37
  layers.0.weight  (4, 2)
  layers.0.bias    (4,)
  layers.1.weight  (4, 4)
  layers.1.bias    (4,)
  layers.2.weight  (1, 4)
  layers.2.bias    (1,)

37 parameters, in the same shapes as the NumPy version. A shorter way to write the same thing is nn.Sequential(nn.Linear(2, 4), nn.Tanh(), nn.Linear(4, 4), nn.Tanh(), nn.Linear(4, 1)). The Lab's generated code does both: a class, PINN(nn.Module), that holds an nn.Sequential of this kind. We use the class form here because it is the one you will extend (Chapter 11 adds input transformations).

Proof that it is the same function

Copy the weights PyTorch chose into our NumPy forward from Section 2.2 and compare outputs. If they match, the two are one function written twice.

def forward(params, x):                       # the NumPy forward pass from Section 2.2
    h = np.atleast_2d(x)
    for W, b in params[:-1]:
        h = np.tanh(h @ W.T + b)
    W, b = params[-1]
    return h @ W.T + b

params = [(l.weight.detach().numpy(), l.bias.detach().numpy()) for l in net.layers]

x = torch.rand(5, 2)                          # five random points in the unit square
out_torch = net(x).detach().numpy()
out_numpy = forward(params, x.numpy())
print("max difference:", float(np.abs(out_torch - out_numpy).max()))

max difference: 2.9802322387695312e-08

The difference is at the level of float32 rounding (about $10^{-7}$), so the two agree.

Initialising the weights

PyTorch's Linear draws its initial weights from a uniform distribution scaled by $1/\sqrt{n_\text{in}}$, in the same spirit as our rule in Section 2.2. The Lab's generated code replaces that with the Xavier (or Glorot) scheme of Glorot & Bengio (2010), which chooses the standard deviation from both sides of the layer,

$$\operatorname{std}(W) = \sqrt{\frac{2}{n_\text{in} + n_\text{out}}},$$

and sets all biases to zero. In PyTorch it is one call per layer:

torch.manual_seed(0)
layer = nn.Linear(40, 40)
nn.init.xavier_normal_(layer.weight)
nn.init.zeros_(layer.bias)
print("sample std of weights:", round(layer.weight.std().item(), 3))
print("formula               :", round((2 / (40 + 40)) ** 0.5, 3))

sample std of weights: 0.159
formula               : 0.158

The sample standard deviation matches the formula to within sampling noise of the 1,600 weights drawn.

What requires_grad means

Every entry in net.parameters() is a tensor with requires_grad=True. That flag tells PyTorch to record the chain of operations that produced any result built from the parameter, so that it can later run the chain backwards. It does not do anything yet:

p = next(net.parameters())
print("requires_grad:", p.requires_grad)
print("gradient before any backward pass:", p.grad)

requires_grad: True
gradient before any backward pass: None

The next chapter makes use of it. When we call .backward() on a loss, every parameter gets a .grad, the slope of the loss with respect to it, and gradient descent takes a step against that slope.

Exercises

1. Build MLP([3, 20, 20, 1]). How many parameters does it have? Check your hand count against the code.
2. What would happen to the equivalence check above if the NumPy forward used np.maximum(0, ·) instead of np.tanh?
3. Why do we call .detach() before .numpy()?

Recap

• A PyTorch network is the same stack of $\sigma(W\mathbf{x} + \mathbf{b})$ layers, stored as an nn.Module with named parameters.
• Defaults differ: PyTorch tensors are float32, NumPy arrays float64. PINN code sometimes switches to float64.
• requires_grad=True is what makes automatic differentiation possible; Chapter 3 uses it for training and Chapter 5 for derivatives with respect to inputs.