PyTorch训练MNIST CNN时出现forward()参数不匹配错误的解决

我用MNIST数据集训练CNN，代码如下：

from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


train_data = datasets.MNIST(
    root="data",
    train=True,
    transform=ToTensor(),
    download=True
    )

test_data = datasets.MNIST(
    root="data",
    train=False,
    transform=ToTensor(),
    download=True
    )

# Process data to batches
loaders = {
    "train": DataLoader(train_data,
                        batch_size=100,
                        shuffle=True,
                        num_workers=0),

    "test": DataLoader(test_data,
                       batch_size=100,
                       shuffle=True,
                       num_workers=0)
}

class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x), 2)))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)

        return F.softmax(x)
    

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = CNN().to(device)

optimizer = optim.Adam(model.parameters(),
                       lr=0.001)

loss_fn = nn.CrossEntropyLoss()


def train(epoch):

    model.train()
    for batch_idx, (data, target) in enumerate(loaders["train"]):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 20 == 0:
            print(f'Train epoch: {epoch} [{batch_idx * len(data)}/{len(loaders["train"].dataset)} ({100. * batch_idx / len(loaders["train"]):.0f}%)]	{loss.item():.6f}')


def test():
    model.eval()

    test_loss = 0
    correct = 0

    with torch.no_grad():
        for data, target in loaders["test"]:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += loss_fn(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(loaders['test'].dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(loaders["test"].dataset)} ({100. * correct / len(loaders["test"].dataset):.0f}%\n)')


for epoch in range(1, 11):
    train(epoch)
    test()

运行时出现以下错误：

Traceback (most recent call last):
  File "c:/Users/u/Projects/venv/src/digits_model_training.py", line 103, in <module>
    train(epoch)
  File "c:/Users/u/Projects/venv/src/digits_model_training.py", line 76, in train
    output = model(data)
  File "C:\Users\u\Projects\venv\venv\lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\u\Projects\venv\venv\lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "c:/Users/u/Projects/venv/src/digits_model_training.py", line 51, in forward
    x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x), 2)))
  File "C:\Users\u\Projects\venv\venv\lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\u\Projects\venv\venv\lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: forward() takes 2 positional arguments but 3 were given

我仅向模型的forward传入了图像数据x，当前用CPU训练，想知道错误原因和修复方法。

错误出在forward函数的这一行：

x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x), 2)))

这里把2作为第二个参数传给了self.conv2_drop（即nn.Dropout2d模块），但Dropout2d.forward()只接受输入张量这一个位置参数，额外传入的2导致参数数量不匹配，触发了TypeError。

实际上，这个2是F.max_pool2d的池化窗口大小参数，错误地放在了conv2_drop的调用参数里，而非max_pool2d的参数位置。

另外还有一个潜在问题：使用nn.CrossEntropyLoss时，模型最后返回了F.softmax(x)。CrossEntropyLoss内部已经包含了log_softmax和nll_loss的计算流程，直接返回全连接层的原始输出（logits）即可，额外做softmax会导致损失计算不准确。

1. 修正forward函数中max_pool2d的参数位置，把2从conv2_drop的调用中移到max_pool2d的参数里。
2. 移除最后的F.softmax，直接返回全连接层的输出。

修正后的CNN类代码如下：

class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        # 修正：把2移到max_pool2d的参数位置
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        # 移除softmax，直接返回logits
        return x

这样修改后，代码就能正常运行，同时损失计算也会更准确。

内容的提问来源于stack exchange，提问作者user26598303