PyTorch训练Pix2pix在Colab P100上触发CUDA非法内存访问错误

核心问题定位 #

报错固定在第33次迭代触发、本地RTX3080Ti可正常跑通、开同步调试后定位到反向传播阶段cuDNN执行失败，本质是Colab默认预装的PyTorch/cuDNN版本和P100显卡存在兼容问题，再叠加代码本身预处理、训练逻辑的隐患共同导致。触发错误后显存没法释放是CUDA上下文损坏的固有表现，不存在占显存的残留进程，除了重启会话没有别的恢复方法。

按优先级排查修复 #

1. 关闭不兼容的硬件加速开关（90%概率直接解决） #

P100是帕斯卡架构的老卡，根本不支持30系、A100以上才有的TF32计算特性，但Colab默认的PyTorch配置是开着TF32开关的，跑到特定尺寸的卷积核的时候直接就触发非法内存访问，和报错里附的官方复现场景完全匹配。
在所有torch相关导入之后、模型定义之前添加以下配置：

import torch
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

如果加完还是报错，可以临时加一行torch.backends.cudnn.enabled = False验证，关闭cuDNN后训练速度会下降，但如果能跑通就能100%确认是cuDNN兼容问题，后续可以把PyTorch降到1.12.0版本彻底解决速度问题。

2. 修复数据预处理逻辑隐患 #

当前代码仅用ToTensor()做转换，没有统一输入尺寸、强制通道数，你设置的batch_size=1，第33次迭代刚好读到第33张样本，要是这张图是RGBA四通道、灰度图、损坏文件、奇数尺寸，就会触发计算阶段的内存错误。
修改transform配置，提前规整输入，避免中间过程产生奇数尺寸的特征图：

transform = transforms.Compose([
    transforms.Lambda(lambda img: img.convert('RGB')), # 强制所有图转3通道RGB
    transforms.Resize((target_shape, target_shape*2)), # 输入图统一为高512、宽1024，切分后左右两半刚好512*512
    transforms.ToTensor(),
])

也可以直接找到训练集里按文件名排序第33张图片，检查是否损坏、格式异常，直接删掉异常样本即可。

3. 删除冗余反向传播参数 #

当前判别器反向传播时加了retain_graph=True，但你训练生成器的时候会重新走完整前向计算，这个参数没有任何作用，只会逐轮累积计算图占显存。P100的显存页管理机制和30系显卡有差异，显存占用累积到阈值就会触发访问错误。
把判别器反向传播那行改成下面这样，删掉冗余参数：

disc_loss.backward()

4. 修正学习率配置 #

当前生成器学习率设为0.01、判别器为0.001，远高于Pix2Pix原论文用的2e-4，学习率过高会导致训练早期参数快速爆炸产生NaN/Inf异常值，传入cuDNN内核时也会触发执行失败。把学习率改成原论文默认值即可：

lr = 2e-4

验证说明 #

按上述顺序修改后，重启Colab会话直接运行即可，不需要改动U-Net和判别器的结构代码，正常情况下不会再出现固定迭代次数报错的问题。

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