求Mac M1芯片上PyTorch的MPS支持torch.nn.Conv3d()的解决方案

问题回顾 #

本科毕设开发图像超分辨率深度学习网络，因使用多页TIFF图像需依赖torch.nn.Conv3d()，但M1芯片的MPS设备不支持该算子。CPU运行速度过慢，此前转用Windows台式机，却因内存不足无法处理高分辨率图像，只能依赖Mac的大内存，同时急需GPU加速来推进调试效率。

可行解决方向 #

1. 升级至PyTorch Nightly版本 #

MPS的算子支持一直在迭代更新，官方Nightly版本大概率已补全Conv3d的MPS支持。直接安装最新Nightly包：

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/mps

安装完成后重启PyCharm，测试Conv3d能否在MPS设备上正常运行——这是最省心的方案，优先尝试。

2. 手动拆分3D卷积为2D卷积组合 #

若Nightly版本仍不支持，可自行将3D卷积拆分为深度维度+空间维度的2D卷积组合，模拟3D卷积效果。示例代码如下：

import torch.nn as nn

class Conv3dTo2d(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, stride=1, padding=0):
        super().__init__()
        # 处理深度维度的分组卷积
        self.depth_conv = nn.Conv2d(in_ch, in_ch, 
                                    kernel_size=(kernel_size[0], 1),
                                    stride=(stride[0], 1),
                                    padding=(padding[0], 0),
                                    groups=in_ch)
        # 处理空间维度的普通2D卷积
        self.spatial_conv = nn.Conv2d(in_ch, out_ch,
                                      kernel_size=(1, kernel_size[1]),
                                      stride=(1, stride[1]),
                                      padding=(0, padding[1]))

    def forward(self, x):
        # 输入shape: (B, C, D, H, W)
        # 合并batch与深度维度，处理深度卷积
        x = x.permute(0, 2, 1, 3, 4).flatten(0, 1)  # (B*D, C, H, W)
        x = self.depth_conv(x)
        # 恢复原维度结构
        x = x.unflatten(0, (-1, x.size(0) // x.size(1))).permute(0, 2, 1, 3, 4)
        # 合并深度与高度维度，处理空间卷积
        x = x.flatten(2, 3)  # (B, C, D*H, W)
        x = self.spatial_conv(x)
        # 恢复3D张量形状
        x = x.unflatten(2, (-1, x.size(2) // x.size(3)))
        return x

需注意对比自定义模块与原生Conv3d的输出结果，确保精度一致，避免影响模型效果。

3. 替换为2D卷积+深度注意力的模型结构 #

多页TIFF的深度维度可视为序列，改用2D卷积处理单页空间信息，再通过注意力机制融合不同页面的深度特征——比如添加Transformer自注意力层或简单门控融合模块，全程使用Conv2d即可完美适配MPS加速，效果未必逊于3D卷积。

4. 混合精度优化（聊胜于无） #

若以上方案均不可行，可尝试自动混合精度，让支持MPS的算子优先跑GPU，不支持的回退CPU，能一定程度缩短运行时间：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler(enabled=True)

# 训练流程示例
with autocast(device_type='mps'):
    output = model(input)
    loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

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