如何修复PyTorch中输入与目标batch_size不匹配的ValueError？

Let's break down the issues in your code and fix them step by step—since you're new to PyTorch, I'll explain each change clearly:

1. Missing Forward Method & Pooling Layers in Your Model #

The biggest culprit here is that your Net class doesn't have a forward pass defined—PyTorch requires this to know how data flows through your layers. Additionally, you're missing pooling layers after your convolutions, which explains why the input dimension to your first fully connected layer (fc1) doesn't align with the actual output shape from your convolutions.

Your fc1 is set to expect 32768 input features (which is 128 * 16 * 16), but without pooling, your convolution layers would output a massive feature map that doesn't match this number. Adding MaxPool2d layers after each convolution downsamples the feature maps to the correct size.

Here's the corrected model definition:

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.pool = nn.MaxPool2d(2, 2)  # Adds pooling to downsample feature maps
        self.conv2 = nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv3 = nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.fc1 = nn.Linear(128 * 16 * 16, 128)  # Matches pooled feature map size
        self.fc2 = nn.Linear(128, 29)
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        # Define how data moves through layers: Conv -> ReLU -> Pool
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        
        # Flatten feature maps for the fully connected layer (-1 infers batch size)
        x = x.view(-1, 128 * 16 * 16)
        
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# Initialize model and move to GPU (if available in Colab)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net().to(device)

2. Fix Training Loop Issues #

a. Batch Size Mismatch Fix #

Without the forward method and pooling layers, your model was outputting a tensor with an incorrect batch dimension (hence the batch_size (1) vs (4) error). The corrected forward pass ensures the output shape is [batch_size, 29] (matching your 29 classes), which aligns with the target batch size.

b. Fix Epoch Printing Bug #

In your original code, you're printing the total number of epochs (epoch) instead of the current training epoch (e). This makes your logs confusing—here's the fix:

# In your print statement, replace `epoch` with `e+1` (to count epochs from 1 instead of 0)
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
    e+1,
    batch_idx * len(data), len(train_loader.dataset),
    100. * batch_idx / len(train_loader), loss.item()))

c. Uncomment Device Transfer #

You commented out the line that moves data to the same device as your model. If you're using Colab's GPU (which you should for faster training), this will cause errors—uncomment it:

data, target = data.to(device), target.to(device)

d. Improve Loss Tracking #

Instead of appending the last batch's loss to your epoch array, track the average loss per epoch for a more accurate training curve:

running_loss = 0.0
# Inside the batch loop:
running_loss += loss.item()
# After the batch loop:
avg_epoch_loss = running_loss / len(train_loader)
loss_epoch_arr.append(avg_epoch_loss)

3. Full Corrected Training Code #

Putting it all together, here's the fixed training function:

import torch.optim as optim
import matplotlib.pyplot as plt

loss_fn = nn.CrossEntropyLoss()
opt = optim.SGD(model.parameters(), lr=0.01)

def train(model, train_loader, optimizer, loss_fn, num_epochs, device):
    model.train()
    loss_epoch_arr = []
    
    for e in range(num_epochs):
        running_loss = 0.0
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            
            output = model(data)
            loss = loss_fn(output, target)
            
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            if batch_idx % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    e+1,
                    batch_idx * len(data), len(train_loader.dataset),
                    100. * batch_idx / len(train_loader), loss.item()))
        
        avg_epoch_loss = running_loss / len(train_loader)
        loss_epoch_arr.append(avg_epoch_loss)
    
    plt.plot(loss_epoch_arr)
    plt.xlabel('Epoch')
    plt.ylabel('Average Loss')
    plt.title('Training Loss Over Epochs')
    plt.show()

# Start training
train(model, trainloader, opt, loss_fn, 10, device)

Why This Works #

• Forward Pass: Defines clear data flow through layers, ensuring the output shape matches your target batch size and class count.
• Pooling Layers: Reduces feature map size to align with your fully connected layer's input dimensions, fixing the batch size mismatch.
• Device Alignment: Ensures both model and data live on the same device (GPU/CPU) to avoid runtime errors.
• Better Loss Tracking: Gives you a more accurate view of how training loss decreases over time.

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