Deep-Learning-with-PyTorch

126 CHAPTER 5 The mechanics of learning
NOTE You might be curious why zeroing the gradient is a required step
instead of zeroing happening automatically whenever we call backward.
Doing it this way provides more flexibility and control when working with gradients
in complicated models.
Having this reminder drilled into our heads, let’s see what our autograd-enabled
training code looks like, start to finish:
# In[9]:
def training_loop(n_epochs, learning_rate, params, t_u, t_c):
for epoch in range(1, n_epochs + 1):
if params.grad is not None:
params.grad.zero_()
t_p = model(t_u, *params)
loss = loss_fn(t_p, t_c)
loss.backward()
with torch.no_grad():
params -= learning_rate * params.grad
if epoch % 500 == 0:
print('Epoch %d, Loss %f' % (epoch, float(loss)))
return params
This could be done at any point in the
loop prior to calling loss.backward().
This is a somewhat cumbersome bit
of code, but as we’ll see in the next
section, it’s not an issue in practice.
Note that our code updating params is not quite as straightforward as we might have
expected. There are two particularities. First, we are encapsulating the update in a
no_grad context using the Python with statement. This means within the with block,
the PyTorch autograd mechanism should look away: 11 that is, not add edges to the forward
graph. In fact, when we are executing this bit of code, the forward graph that
PyTorch records is consumed when we call backward, leaving us with the params leaf
node. But now we want to change this leaf node before we start building a fresh forward
graph on top of it. While this use case is usually wrapped inside the optimizers
we discuss in section 5.5.2, we will take a closer look when we see another common use
of no_grad in section 5.5.4.
Second, we update params in place. This means we keep the same params tensor
around but subtract our update from it. When using autograd, we usually avoid inplace
updates because PyTorch’s autograd engine might need the values we would be
modifying for the backward pass. Here, however, we are operating without autograd,
and it is beneficial to keep the params tensor. Not replacing the parameters by assigning
new tensors to their variable name will become crucial when we register our
parameters with the optimizer in section 5.5.2.
11 In reality, it will track that something changed params using an in-place operation.