CPU可运行的MNIST卷积网络在MPS设备报NNPack张量类型不匹配错误

可能的核心原因 #

• 数据类型不匹配：MPS后端对张量数据类型的兼容性要求更严格，若输入数据与模型参数的 dtype 不一致（比如一个是 float32，一个是 float64），会触发卷积层的类型匹配错误。
• Softmax维度设置错误：你的代码中Softmax(dim=0)是对 batch 维度做归一化，但MNIST任务中应该对每个样本的类别维度（dim=1）做Softmax。这种逻辑错误在CPU上可能未触发异常，但MPS后端的张量处理逻辑会间接引发类型不匹配报错。
• NNPACK与MPS的兼容冲突：NNPACK是CPU专用的卷积加速库，当模型切换到MPS时，若PyTorch仍尝试调用NNPACK的卷积实现，会因设备不兼容导致类型错误。

具体解决步骤 #

1. 统一张量数据类型 #

确保输入数据和模型参数都使用float32（MPS对float32支持最优）：

• 加载MNIST数据时指定dtype：

  transform = transforms.Compose([
      transforms.ToTensor(),
      transforms.Lambda(lambda x: x.to(torch.float32))
  ])
  train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
  

• 模型转到MPS时确认dtype：

  model = NeuralNetwork().to('mps', dtype=torch.float32)
  

2. 修正Softmax的维度 #

将nn.Softmax(dim=0)改为nn.Softmax(dim=1)，确保每个样本的10个类别输出做归一化：

self.mnist_nn = nn.Sequential(
    # ... 其他层保持不变
    nn.Linear(128, 10),
    nn.Softmax(dim=1)
)

3. 强制使用MPS卷积实现 #

若仍存在NNPACK兼容问题，可通过设置环境变量强制PyTorch使用MPS的卷积后端：

export PYTORCH_ENABLE_MPS_FALLBACK=1

或者在代码开头添加：

import os
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'

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