machine-learning - 损失函数如何知道在 PyTorch 中为哪个模型计算梯度？

Question

我不确定 PyTorch 如何将损失函数链接到我想要计算的模型。损失和模型之间从来不存在明确的引用，例如模型参数和优化器之间的引用。

例如，我想在同一数据集上训练 2 个网络，因此我想利用一次数据集传递。 PyTorch 如何将适当的损失函数链接到适当的模型。引用代码如下:

import torch
from torch import nn, optim
import torch.nn.functional as F
from torchvision import datasets, transforms
import shap

# Define a transform to normalize the data
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,)),
                              ])
# Download and load the training data
trainset = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

model = nn.Sequential(nn.Linear(784, 128),
                      nn.ReLU(),
                      nn.Linear(128, 64),
                      nn.ReLU(),
                      nn.Linear(64, 10),
                      nn.LogSoftmax(dim=1))

model2 = nn.Sequential(nn.Linear(784, 128),
                      nn.ReLU(),
                      nn.Linear(128, 10),
                      nn.LogSoftmax(dim=1))

# Define the loss
criterion = nn.NLLLoss()
criterion2 = nn.NLLLoss()

optimizer = optim.SGD(model.parameters(), lr=0.003)
optimizer2 = optim.SGD(model2.parameters(), lr=0.003)

epochs = 5
for e in range(epochs):
    running_loss = 0
    running_loss_2 = 0
    for images, labels in trainloader:
        # Flatten MNIST images into a 784 long vector
        images = images.view(images.shape[0], -1) # batch_size x total_pixels

        # Training pass
        optimizer.zero_grad()
        optimizer2.zero_grad()

        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()


        output2 = model2(images)
        loss2 = criterion2(output2, labels)
        loss2.backward()
        optimizer2.step()

        running_loss += loss.item()
        running_loss_2 += loss2.item()

    print(f"Training loss 1: {running_loss/len(trainloader)}")
    print(f"Training loss 2: {running_loss_2/len(trainloader)}")
    print()

那么，当调用 loss.backward() 和 loss2.backward() 时，pytorch 如何知道如何计算适当模型的适当梯度？

Answer 1

每当您使用模型参数之一(或具有属性 torch.tensor 的任何 requires_grad==True )执行前向操作时，pytorch 都会构建一个计算图。当您对该图中的后代进行操作时，该图将被扩展。就您而言，您有一个 nn.module叫model其中会有一些可训练的model.parameters() ，因此 pytorch 将从您的 model.parameters() 构建一个图表当您执行远期操作时，会一直导致损失。然后在向后传递期间反向遍历该图，以将梯度传播回参数。对于 loss在上面的代码中，图表类似于

model.parameters() --> [intermediate variables in model] -->  output --> loss
                                  ^                                        ^
                                  |                                        |
                               images                                     labels

当您调用loss.backward()时pytorch 反向遍历该图以达到所有可训练参数(本例中仅 model.parameters())并更新 param.grad对于他们每个人来说。 optimizer然后依靠在向后传递过程中收集的信息来更新参数。 对于 loss2故事很相似。

官方pytorch tutorials是有关这方面更深入信息的良好资源。