PyTorch自编码器模型并行实现遇CUDA OOM及设备不匹配问题咨询

报错的根本原因是你在Autoencoder构造函数中已经分别将encoder部署到cuda:0、decoder部署到cuda:1，但后续初始化模型时调用了.to(device)，这个操作会把模型所有参数统一迁移到device指定的单张卡上，直接覆盖了构造函数里的设备分配逻辑，导致decoder的实际参数所在设备和forward中x迁移到的cuda:1不一致，触发矩阵运算的设备匹配错误。
解决步骤如下：

• 移除模型初始化时的.to(device)调用，修改为model = Autoencoder(input_output_size)即可，保留构造函数内的模块设备分配逻辑。
• 调整损失计算的张量设备匹配规则：当前输入images被迁移到了device，模型输出默认在cuda:1，你可以选择将images迁移到cuda:1计算损失，避免输出张量的多余迁移，降低设备通信开销：train_loss = criterion(output_images, images.to('cuda:1'))
• 你当前的优化器初始化逻辑是对的，只要删除模型初始化时的.to(device)，优化器会自动适配不同设备上的参数。

你当前单卡12GB依然OOM的核心原因是用全连接层构建自编码器：3×256×256的图像展平后是196608维，第一层和最后一层全连接的单参数量就超过2亿，加上Adam优化器的状态、梯度、中间激活值的占用，即使batch size为1也会超出显存上限。优化方案按优先级排序如下：

• 优先替换为卷积自编码器结构：输入输出直接用3维张量（不需要展平），编码器用堆叠的nn.Conv2d+nn.ReLU+nn.MaxPool2d（通过padding保证特征图尺寸匹配），解码器用nn.ConvTranspose2d或nn.Upsample+nn.Conv2d恢复到256×256尺寸，参数量会降低2个数量级以上，12GB单卡batch size可以开到16以上，不需要模型并行即可正常训练。
• 如果必须保留全连接结构，可采用以下显存优化手段：
  • 开启混合精度训练：使用PyTorch自带的AMP模块，用半精度存储参数和计算，显存占用直接降低50%左右，同时训练速度也会提升。
  • 启用梯度检查点：用torch.utils.checkpoint.checkpoint包装编码器、解码器的层，反向传播时重新计算中间激活值而非提前存储，可降低30%~60%的显存占用，仅损失少量训练速度。
  • 替换优化器：将Adam替换为SGD，SGD不需要存储每个参数的两个动量状态，可减少约33%的参数相关显存占用。
  • 移除冗余操作：删除每个batch末尾的gc.collect()和torch.cuda.empty_cache()调用，这两个操作不会降低峰值显存占用，反而会大幅拖慢训练速度，仅需在每个epoch结束时按需调用即可。
  • 细粒度模型并行：将参数量最大的第一层和最后一层全连接拆分为多个子层，分别部署到两张GPU上，进一步降低单卡的显存压力。

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