基于PyTorch-CNN的图像处理高损失问题求助

基于PyTorch搭建的CNN图像处理代码可正常运行，但损失值高达21%。因设备算力限制，已将原本600x600的输入图像调整为256x256，并在训练轮次中转为灰度图。尝试训练10个epochs，损失值并未下降，恳请提供排查建议或指出代码问题。

部分代码 #

lr = 1e-3  
batch_size = 1
num_epochs = 1
inputs=[np.array(Image.fromarray(a).resize((256,256))) for a in inputs]
labels=[np.array(Image.fromarray(a).resize((256,256))) for a in labels]
inputs_tensor = torch.Tensor(inputs).cuda()    # convert to tensor
labels_tensor = torch.Tensor(labels).cuda()
print(inputs_tensor.shape)

...

import sys
import numpy as np
# Testing loop
for epoch in range(num_epochs):
    
    total_loss = 0

    model.eval()
    
    total_test_loss = 0

    correct_predictions = 0
    total_samples = 0
    total_samples_predict = 0

    # Store the input, ground truth, and predicted images for visualization
    input_images = []
    ground_truth_images = []
    predicted_images = []

    for delayed_tec_batch, label_batch in test_loader:
        with torch.no_grad():
            output = model(delayed_tec_batch)
            test_loss = criterion(output, label_batch)  # Flatten label
            total_test_loss += test_loss.item()
            
            output = output.permute(0, 2, 3, 1)
            label_batch = label_batch.permute(0, 2, 3, 1)
            
            #print ("output Shape ", output.shape) #predicted values
            #print ("label batch Shape ", label_batch.shape) #ground truth
            
            # Check if output is a numpy array, if not, convert it
            if isinstance(output, torch.Tensor) and output.is_cuda:
                output = output.cpu()
            output = np.array(output)
            
            
            if isinstance(label_batch, torch.Tensor) and label_batch.is_cuda:
                label_batch = label_batch.cpu()
            label_batch = np.array(label_batch)

                
            output_gray = cv2.cvtColor(output[0], cv2.COLOR_BGR2GRAY)
            (_, output_gray2) = cv2.threshold(output_gray, 1, 255, cv2.THRESH_BINARY)
            
            label_batch_gray = cv2.cvtColor(label_batch[0], cv2.COLOR_BGR2GRAY)
            (_, label_batch_gray2) = cv2.threshold(label_batch_gray, 1, 255, cv2.THRESH_BINARY)
            
            correct_predictions += np.sum((output_gray2 == label_batch_gray2 ))
            total_samples += np.sum(label_batch_gray2 == 255)


    average_test_loss = total_test_loss / len(test_loader)
    accuracy = correct_predictions / total_samples

    print(f"Testing Loss: {average_test_loss:.4f}")
    print(f"Accuracy: {accuracy:7f}%")
    
    print ("correct_predictions ", correct_predictions)
    print ("total_samples  ", total_samples)
    print ("total_samples_predict  ", total_samples_predict)
    print ("output ", output_gray2)
    print ("label batch ", label_batch_gray2)
    
    
    # Convert tensors to numpy arrays for visualization
    input_images.extend(delayed_tec_batch)#.cpu().numpy())
    ground_truth_images.extend(label_batch)#.cpu().numpy())
    predicted_images.extend(output)#.cpu().numpy())
 
    # Plot the learning curve
    plt.plot(total_test_loss , label='Training Loss')
    plt.plot(validation_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Learning Curve')
    plt.legend()
    plt.show()

输出结果 #

torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
torch.Size([1, 64, 4, 4])
**Testing Loss: 21.7260**
**Accuracy: 5.137234%**
**correct_predictions  4456910**
**total_samples   867570**
total_samples_predict   0
output  [[255. 255. 255. ... 255. 255. 255.]
 [255.   0.   0. ...   0.   0. 255.]
 [255.   0.   0. ...   0.   0. 255.]
 ...
 [255.   0.   0. ...   0.   0. 255.]
 [255.   0.   0. ...   0.   0. 255.]
 [255. 255. 255. ... 255. 255. 255.]]
label batch  [[255. 255. 255. ... 255. 255. 255.]
 [255. 255. 255. ... 255. 255. 255.]
 [255.   0. 255. ...   0.   0. 255.]
 ...
 [255.   0.   0. ...   0.   0. 255.]
 [255. 255. 255. ... 255. 255. 255.]
 [255. 255. 255. ... 255. 255. 255.]]

排查建议 #

• 核心问题：缺少训练循环：当前代码仅展示了测试逻辑，没有model.train()模式切换、梯度反向传播(loss.backward())和优化器更新(optimizer.step())的代码，模型参数从未更新，这是损失不下降的根本原因。
• 数据预处理错误：
  • 标签图像resize时应使用Image.NEAREST插值，避免插值引入非0/255的中间值，干扰后续阈值处理；
  • 灰度转换应在数据加载阶段完成，确保模型输入是单通道张量，而非在测试阶段才转灰度，否则模型输入通道数与任务需求不匹配；
  • 输出与标签维度不匹配：从打印的torch.Size([1,64,4,4])来看，模型输出是64通道、4x4尺寸的特征图，但标签是256x256的掩码，尺寸和通道数均不匹配，导致损失计算完全错误。
• 损失函数选型不当：若为图像分割任务，MSE损失不适合，建议改用DiceLoss、BCEWithLogitsLoss（二分类）或CrossEntropyLoss（多分类），并确保输出经过激活函数后与标签格式一致。
• 评估指标计算错误：
  • 准确率计算逻辑错误：correct_predictions统计所有匹配像素，total_samples仅统计标签中255的像素，这样的准确率无参考价值，应使用IoU、Dice系数等分割专用指标，或用所有像素作为分母；
  • total_samples_predict未赋值，属于无效变量，建议删除或补充计算逻辑。
• 超参数不合理：
  • batch_size=1会导致梯度噪声过大，建议在算力允许范围内调大（如4、8）；
  • 初始学习率1e-3过高，可尝试1e-4，并配合ReduceLROnPlateau学习率调度器动态调整；
  • 训练轮次需结合任务难度调整，同时必须在训练循环中更新参数，否则增加epochs无效。
• 模型结构问题：若为图像分割任务，模型需要添加上采样层（转置卷积、Upsample等），将64x4x4的特征图恢复到256x256的输入尺寸，保证输出与标签尺寸一致。

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