PyTorch人脸关键点检测CNN模型通道不匹配RuntimeError问题

错误原因分析 #

你遇到的RuntimeError: Given groups=1, weight of size [32, 1, 3, 3], expected input[1, 8, 128, 128] to have 1 channels, but got 8 channels instead错误，核心原因是训练函数中错误地移除了图像的通道维度，导致模型接收到的输入形状不符合预期：

• 模型第一层卷积conv1定义为输入1通道，但train函数里的images = images.squeeze(1)把原本形状为[8,1,128,128]（批量大小8、1通道、128x128分辨率）的张量，压缩成了[8,128,128]
• 模型会自动将这个形状解读为[批量大小=1, 通道数=8, 高=128, 宽=128]，完全不符合设计，从而触发通道不匹配错误。

同时代码还存在其他潜在问题（缺失导入、全连接层维度计算错误等），下面是完整的修正方案：

具体修正步骤 #

1. 补充缺失的导入语句 #

代码中缺少多个必要的库导入，需添加：

import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from PIL import Image
import torch.optim as optim

2. 移除训练函数中的错误维度操作 #

删除train函数里的images = images.squeeze(1)，保留图像的通道维度；同时需将标注的关键点展平为一维，匹配模型输出的136维度：

def train(model, train_loader, criterion, optimizer, device):
    model.train()
    train_loss = 0.0
    for images, keypoints in train_loader:
        images, keypoints = images.to(device), keypoints.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        # 将keypoints展平为一维，匹配模型输出
        loss = criterion(outputs, keypoints.view(keypoints.size(0), -1))
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * images.size(0)
        
    return train_loss / len(train_loader.dataset)

3. 修正评估函数的维度操作 #

删除evaluate函数里的inputs = inputs.unsqueeze(1).float()，直接使用原始输入，同时展平标注维度：

def evaluate(model, val_loader, criterion, device):
    model.eval()
    running_loss = 0.0
    with torch.no_grad():
        for data in val_loader:
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.view(labels.size(0), -1).float().to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item()
    return running_loss / len(val_loader)

4. 修正全连接层的输入维度 #

经过4次MaxPool2d(stride=2)，输入图像128x128会依次缩小为64x64 → 32x32 → 16x16 → 8x8，因此全连接层fc1的输入维度需修正：

self.fc1 = nn.Linear(256 * 8 * 8, 1024)

5. 修复测试集加载遗漏 #

主函数中定义了test_set但未创建test_loader，需添加：

test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)

修正后的完整代码 #

import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision.transforms import ToTensor, transforms
from PIL import Image

class FaceKeypointDataset(Dataset):
    def __init__(self, image_folder_path, keypoint_folder_path, transform=None):
        self.image_folder_path = image_folder_path
        self.keypoint_folder_path = keypoint_folder_path
        self.transform = transform
        # 确保图像和标注文件名一一对应
        self.image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_folder_path) if f.endswith('.png') or f.endswith('.jpg')]
        # 过滤出和图像匹配的标注文件
        self.keypoint_filenames = [f for f in self.image_filenames if os.path.exists(os.path.join(keypoint_folder_path, f + '.txt'))]

    def __len__(self):
        return len(self.keypoint_filenames)

    def __getitem__(self, idx):
        filename = self.keypoint_filenames[idx]
        image_path = os.path.join(self.image_folder_path, filename + '.png')
        keypoint_path = os.path.join(self.keypoint_folder_path, filename + '.txt')
        
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        keypoint = np.loadtxt(keypoint_path, delimiter=',')
        # Normalize keypoint coordinates to [-1, 1]
        keypoint[:, 0] = (keypoint[:, 0] / image.shape[1]) * 2 - 1
        keypoint[:, 1] = (keypoint[:, 1] / image.shape[0]) * 2 - 1

        if self.transform:
            image = Image.fromarray(image) # convert numpy ndarray to PIL Image
            image = self.transform(image)

        return image, torch.from_numpy(keypoint).float()

transform = transforms.Compose([transforms.Resize(128),
                                transforms.ToTensor()]) # 已经是灰度图，无需再转灰度

class FaceKeypointCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 修正全连接层输入维度：128经过4次池化后是8x8
        self.fc1 = nn.Linear(256 * 8 * 8, 1024)
        self.fc2 = nn.Linear(1024, 136) # 68个关键点，每个x/y，共136维

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = F.relu(self.conv3(x))
        x = self.pool3(x)
        x = F.relu(self.conv4(x))
        x = self.pool4(x)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)

        return x

def train(model, train_loader, criterion, optimizer, device):
    model.train()
    train_loss = 0.0
    for images, keypoints in train_loader:
        images, keypoints = images.to(device), keypoints.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        # 将keypoints展平为一维，匹配模型输出
        loss = criterion(outputs, keypoints.view(keypoints.size(0), -1))
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * images.size(0)
        
    return train_loss / len(train_loader.dataset)

def evaluate(model, val_loader, criterion, device):
    model.eval()
    running_loss = 0.0
    with torch.no_grad():
        for data in val_loader:
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.view(labels.size(0), -1).float().to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item()
    return running_loss / len(val_loader)

def save_best_model(model, val_loss, best_val_loss, model_save_path):
    if val_loss < best_val_loss:
        torch.save(model.state_dict(), model_save_path)
        print(f"Saved best model with validation loss {val_loss:.4f}")
        return val_loss
    return best_val_loss

if __name__ == "__main__":
    # Define hyperparameters
    image_folder_path = "/content/drive/MyDrive/kaggle/Task 1/ann_done"
    keypoint_folder_path = "/content/drive/MyDrive/kaggle/Task 1/ann"
    model_save_path = "/content/drive/MyDrive/kaggle/Task 1/model.pth"
    batch_size = 8
    learning_rate = 0.001
    num_epochs = 20
    test_ratio = 0.2
    val_ratio = 0.2
    random_seed = 42

    # Define device (GPU if available, else CPU)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Define dataset and transforms
    dataset = FaceKeypointDataset(image_folder_path, keypoint_folder_path, transform=transform)

    # Split dataset into train, validation, and test sets
    num_data = len(dataset)
    num_test = int(num_data * test_ratio)
    num_val = int(num_data * val_ratio)
    num_train = num_data - num_test - num_val
    train_set, val_set, test_set = random_split(dataset, [num_train, num_val, num_test],
                                                 generator=torch.Generator().manual_seed(random_seed))

    # Define data loaders
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False) # 新增测试集加载

    # Define model, loss function, and optimizer
    model = FaceKeypointCNN().to(device)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # Train model and save best model based on validation loss
    best_val_loss = np.inf
    for epoch in range(num_epochs):
        train_loss = train(model, train_loader, criterion, optimizer, device)
        val_loss = evaluate(model, val_loader, criterion, device)
        print(f"Epoch {epoch+1}/{num_epochs}, train loss: {train_loss:.4f}, val loss: {val_loss:.4f}")
        best_val_loss = save_best_model(model, val_loss, best_val_loss, model_save_path)

    # Load best model and evaluate on test set
    model.load_state_dict(torch.load(model_save_path))
    test_loss = evaluate(model, test_loader, criterion, device)
    print(f"Test loss: {test_loss:.4f}")

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