PyTorch CNN张量设备不匹配错误求助：定位GPU迁移位置

运行PyTorch搭建的CNN模型时触发如下错误：

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument target in method wrapper_nll_loss_forward)

无法定位需迁移至GPU的代码位置，虽怀疑问题出在损失函数环节，但调整后仍未解决。

关键代码片段 #

数据预处理 #

transformer = transforms.Compose([
    transforms.Resize((350,350)),
    transforms.ToTensor(), 
    transforms.Normalize([0.5,0.5,0.5],
                         [0.5,0.5,0.5])
])

模型定义 #

class ConvNet(nn.Module):
    def __init__(self,num_classes=4):
        super(ConvNet,self).__init__()
        
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=128,kernel_size=3,stride=1,padding='valid')
        self.bn1 = nn.BatchNorm2d(num_features=128)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        
        self.conv2 = nn.Conv2d(in_channels=128,out_channels=64,kernel_size=3,stride=1,padding='valid')
        self.bn2 = nn.BatchNorm2d(num_features=64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2)
        
        self.conv3 = nn.Conv2d(in_channels=64,out_channels=64,kernel_size=3,stride=1,padding='valid')
        self.bn3 = nn.BatchNorm2d(num_features=64)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2)
        
        self.conv4 = nn.Conv2d(in_channels=64,out_channels=32,kernel_size=3,stride=1,padding='valid')
        self.bn4 = nn.BatchNorm2d(num_features=32)
        self.relu4 = nn.ReLU()
        self.pool4 = nn.MaxPool2d(kernel_size=2)
        
        self.conv5 = nn.Conv2d(in_channels=32,out_channels=32,kernel_size=3,stride=1,padding='valid')
        self.bn5 = nn.BatchNorm2d(num_features=32)
        self.relu5 = nn.ReLU()
        self.pool5 = nn.MaxPool2d(kernel_size=2)
        
        self.flat = nn.Flatten()
        self.fc1 = nn.Linear(in_features=2592, out_features = 256)
        self.fc2 = nn.Linear(in_features=256, out_features = num_classes)

        
    def forward(self,input):
        output = self.conv1(input)
        output = self.bn1(output)
        output = self.relu1(output)
        output = self.pool1(output)
        
        output = self.conv2(output)
        output = self.bn2(output)
        output = self.relu2(output)
        output = self.pool2(output)
        
        output = self.conv3(output)
        output = self.bn3(output)
        output = self.relu3(output)
        output = self.pool3(output)
        
        output = self.conv4(output)
        output = self.bn4(output)
        output = self.relu4(output)
        output = self.pool4(output)
        
        output = self.conv5(output)
        output = self.bn5(output)
        output = self.relu5(output)
        output = self.pool5(output)
        
        output = self.flat(output)
        output = self.fc1(output)
        output = self.fc2(output)
        
        return output

训练初始化与循环 #

model = ConvNet(num_classes=4).to(device)
optimizer = Adam(model.parameters(),lr=0.001,weight_decay=0.0001)
loss_function = nn.CrossEntropyLoss()

best_accuracy = 0.0
for epoch in range(num_epochs):
    
    model.train()
    train_accuracy = 0.0
    train_loss = 0.0
    
    for i, (images,labels) in enumerate(train_loader):
        if torch.cuda.is_available():
            images = Variable(images.cuda())
            lables = Variable(labels.cuda())
    
        optimizer.zero_grad()
        outputs = model(images)
        loss = loss_function(outputs,labels)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.cpu().data*images.size(0)
        _.prediction = torch.max(outputs.data,1)
        train_accuracy += int(torch.sum(prediction==labels.data))
    
    train_accuracy = train_accuracy/train_count
    train_loss = train_loss/train_count
    
    # 测试集评估
    model.eval()
    test_accuracy = 0.0
    for i, (images,labels) in enumerate(train_loader):
        if torch.cuda.is_available():
            images = Variable(images.cuda())
            lables = Variable(labels.cuda())
        
        outputs = model(images)
        _.prediction = torch.max(outputs.data,1)
        
    test_accuracy = test_accuracy/test_count

    print('Epoch: '+str(epoch)+' Train Loss: '+str(int(train_loss)))+' Train Accuracy: '+str(train_accuracy)+' Test Accuracy: '+str(test_accuracy)

    if test_accuracy > best_accuracy:
        torch.save(model.state_dict(), 'best_checkpoint.model')

核心问题与修复方案 #

错误根源 #

1. 变量拼写错误：将标签移至GPU时，把labels写成了lables，导致原始labels仍留在CPU，与GPU上的模型输出计算损失时触发设备不匹配
2. 张量解构错误：_.prediction = torch.max(outputs.data,1)写法错误，无法正确获取预测结果
3. 测试集遍历错误：测试循环遍历的是train_loader而非test_loader，且未累加测试准确率

修复后的训练循环代码 #

# 提前统一设备定义
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

best_accuracy = 0.0
for epoch in range(num_epochs):
    
    model.train()
    train_accuracy = 0.0
    train_loss = 0.0
    
    for i, (images,labels) in enumerate(train_loader):
        # 统一迁移到指定设备，无需手动封装Variable
        images = images.to(device)
        labels = labels.to(device)
    
        optimizer.zero_grad()
        outputs = model(images)
        loss = loss_function(outputs,labels)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item() * images.size(0)
        # 正确解构torch.max的返回值
        _, prediction = torch.max(outputs.data, 1)
        train_accuracy += int(torch.sum(prediction == labels.data))
    
    train_accuracy = train_accuracy / train_count
    train_loss = train_loss / train_count
    
    # 测试集评估：遍历test_loader
    model.eval()
    test_accuracy = 0.0
    with torch.no_grad():
        for i, (images,labels) in enumerate(test_loader):
            images = images.to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            _, prediction = torch.max(outputs.data, 1)
            test_accuracy += int(torch.sum(prediction == labels.data))
            
    test_accuracy = test_accuracy / test_count

    # 使用f-string简化打印
    print(f'Epoch: {epoch} Train Loss: {int(train_loss)} Train Accuracy: {train_accuracy:.4f} Test Accuracy: {test_accuracy:.4f}')

    if test_accuracy > best_accuracy:
        torch.save(model.state_dict(), 'best_checkpoint.model')
        best_accuracy = test_accuracy

额外优化建议 #

• PyTorch 0.4+版本无需手动使用Variable，张量默认支持梯度追踪
• 统一使用device变量管理设备，避免多处写.cuda()，提升代码兼容性
• 测试循环添加torch.no_grad()禁用梯度计算，节省显存
• 使用loss.item()替代loss.cpu().data，更简洁符合最佳实践

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