损失函数如何知道在 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()

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

Answer 1

每当您使用模型参数之一（或任何具有属性）执行前向操作时，pytorch 都会构建一个计算图。当您对该图中的后代进行操作时，该图将被扩展。你的情况，你有一个叫这将有一些可训练的，所以pytorch将建立从图表的损失一路为您进行着操作。然后在向后传递期间反向遍历图以将梯度传播回参数。因为 在上面的代码中，图表类似于torch.tensorrequires_grad==Truenn.modulemodelmodel.parameters()model.parameters()loss

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

当您调用loss.backward()pytorch 时，会反向遍历此图以获取所有可训练参数（仅model.parameters()在本例中）并param.grad为每个参数进行更新。在optimizer随后依靠向通行更新参数过程中收集这些信息。因为loss2故事很相似。

官方pytorch 教程是一个很好的资源，可以提供更深入的信息。