PyTorch用CUDA训练ResNet18卡顿且性能劣于CPU，求问题排查

架构老林

2026-7-25

问题描述

我尝试微调ResNet18的最后一层，但使用CUDA时遇到问题。GPU无运行迹象，任务管理器中GPU使用率极低。我将每张图像的张量数增加到5，原本预期会影响性能，但没想到会严重到一整晚仍未完成第一个epoch的程度。

我大致遵循PyTorch官方迁移学习教程。补充信息：数据集包含11个图像类别，多数类别有3000张图像，部分类别仅有几百张。

我已检查驱动与CUDA兼容性，确认CUDA 12.1与当前驱动版本兼容，显卡算力为7.5也在支持范围内。我怀疑是PyTorch的使用方式存在问题，但不确定具体原因。

显卡为GTX-1650，我认为CUDA/GPU根本没有工作，因为尝试可视化图像时程序会完全冻结。

数据导入与张量预处理

def stack_tensor(crops):
    return torch.stack([transforms.ToTensor()(crop) for crop in crops])

def normalize_tensor(crops):
    return torch.stack([transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(crop) for crop in crops])

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize(512),
        transforms.FiveCrop(224),
        transforms.Lambda(stack_tensor),
        transforms.Lambda(normalize_tensor)
    ]),
    'val': transforms.Compose([
        transforms.Resize(512),
        transforms.FiveCrop(224),
        transforms.Lambda(stack_tensor),
        transforms.Lambda(normalize_tensor)
    ]),
}

data_dir = ''
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

模型训练代码

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

训练调用代码（基于CPU预训练权重继续训练）

# Create new model
model_ft = models.resnet18(weights=ResNet18_Weights.DEFAULT)

# Load existing model
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 11)
model_ft.load_state_dict(torch.load("initialmodel"))

# num_ftrs = model_ft.fc.in_features
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
# model_ft.fc = nn.Linear(num_ftrs, 11)

model_ft = model_ft.to(device)

criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=7)

问题分析与解决办法

核心问题1：FiveCrop后的张量维度不匹配，导致模型计算效率极低

你的数据处理中，FiveCrop返回5张裁剪图像，经stack_tensor后输入张量维度为[batch_size, 5, 3, 224, 224]，但ResNet18默认接受[batch_size, 3, 224, 224]的输入。直接喂入5维张量会让PyTorch自动合并维度，不仅损失计算和样本统计逻辑完全错误，还会导致GPU无法高效利用，大部分时间卡在维度适配或CPU预处理环节。

解决办法（二选一）：

方法一：将裁剪样本展平到batch维度，匹配标签后计算损失
修改数据预处理和训练循环：

# 优化预处理，直接生成可展平的张量
def process_five_crop(crops):
    tensors = torch.stack([transforms.ToTensor()(crop) for crop in crops])
    return transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(tensors)

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize(512),
        transforms.FiveCrop(224),
        transforms.Lambda(process_five_crop),
    ]),
    'val': transforms.Compose([
        transforms.Resize(512),
        transforms.FiveCrop(224),
        transforms.Lambda(process_five_crop),
    ]),
}

# 训练循环中调整标签和统计逻辑
for inputs, labels in dataloaders[phase]:
    batch_size, num_crops, c, h, w = inputs.shape
    # 展平为[batch*5, c, h, w]
    inputs = inputs.view(-1, c, h, w).to(device)
    # 重复标签匹配裁剪样本数
    labels = labels.repeat_interleave(num_crops).to(device)

    optimizer.zero_grad()
    with torch.set_grad_enabled(phase == 'train'):
        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)

        if phase == 'train':
            loss.backward()
            optimizer.step()

    # 统计时还原原始样本数
    running_loss += loss.item() * batch_size
    running_corrects += torch.sum(preds == labels.data) // num_crops

方法二：对裁剪样本的输出取平均，再计算损失
更贴合FiveCrop的评估逻辑，训练时也适用：

# 训练循环中处理5维输入
for inputs, labels in dataloaders[phase]:
    batch_size, num_crops, c, h, w = inputs.shape
    inputs = inputs.view(-1, c, h, w).to(device)
    labels = labels.to(device)

    optimizer.zero_grad()
    with torch.set_grad_enabled(phase == 'train'):
        outputs = model(inputs)
        # 重新reshape后取平均输出
        outputs = outputs.view(batch_size, num_crops, -1).mean(dim=1)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)

        if phase == 'train':
            loss.backward()
            optimizer.step()

    running_loss += loss.item() * batch_size
    running_corrects += torch.sum(preds == labels.data)

核心问题2：预处理代码效率低下

你的stack_tensor和normalize_tensor用CPU循环处理每个裁剪图像，耗时极长。可以合并为批量操作优化：

def process_five_crop(crops):
    tensors = torch.stack([transforms.ToTensor()(crop) for crop in crops])
    # 批量归一化，无需循环处理每个张量
    return transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(tensors)

其他排查点

确认PyTorch CUDA可用性：在代码开头添加print(torch.cuda.is_available())和print(device)，确保输出True和cuda:0。
调整数据加载线程：如果CPU性能不足，num_workers=4可能成为瓶颈，尝试改为2或0测试。
验证模型设备：加载权重后添加print(next(model_ft.parameters()).device)，确认模型参数已移至GPU。

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

PyTorch用CUDA训练ResNet18卡顿且性能劣于CPU，求问题排查

数据导入与张量预处理 #

模型训练代码 #

训练调用代码（基于CPU预训练权重继续训练） #

核心问题1：FiveCrop后的张量维度不匹配，导致模型计算效率极低 #

核心问题2：预处理代码效率低下 #