MAC M1 MPS环境下PyTorch训练模型报aten::hardswish_未实现错误

问题背景 #

在M1芯片的Mac上使用MPS后端训练基于MobileNetV3的FasterRCNN实例分割模型时，触发NotImplementedError: The operator 'aten::hardswish_' is not currently implemented for the MPS device错误，CPU环境下模型可正常运行，PyTorch版本为1.13.1，MPS已确认可用。

可行解决方案 #

1. 替换为MPS支持更好的Backbone #

MobileNetV3的Hardswish算子在PyTorch 1.13.1的MPS后端未实现，直接换回ResNet50 backbone即可避开这个问题，修改模型定义代码：

def get_model_instance_segmentation(num_classes):
    # 改用ResNet50的FasterRCNN，MPS对ResNet系列算子支持更完善
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

2. 自定义MPS兼容的Hardswish实现 #

手动替换模型中的Hardswish层为用MPS支持的算子组合实现的版本，示例代码如下：

import torch
import torch.nn as nn

class MPSCompatibleHardswish(nn.Module):
    def forward(self, x):
        # 用torch.clamp和torch.where模拟Hardswish，替代原生的hardswish_算子
        return x * torch.where(x >= 3, torch.tensor(1.0, device=x.device),
                               torch.where(x <= -3, torch.tensor(0.0, device=x.device),
                                           (x + 3) / 6))

# 递归替换模型中的所有Hardswish层
def replace_hardswish(model):
    for name, module in model.named_children():
        if isinstance(module, nn.Hardswish):
            setattr(model, name, MPSCompatibleHardswish())
        else:
            replace_hardswish(module)

# 初始化模型后执行替换
model = get_model_instance_segmentation(config.num_classes)
replace_hardswish(model)
if torch.backends.mps.is_available():
    model.to(device)

3. 升级PyTorch到2.0+版本 #

PyTorch 2.0及后续版本已经完善了对MPS后端的算子支持，其中包括Hardswish算子。通过conda升级PyTorch：

conda install pytorch torchvision torchaudio -c pytorch

升级后重新运行训练代码即可。

4. 部分算子强制回退到CPU #

如果不想修改模型或升级版本，可以将包含Hardswish的层强制放在CPU上运行，示例如下：

# 定位到backbone中包含Hardswish的子模块，将其移到CPU
model.backbone.body.backbone[-1].conv[-1].activation = model.backbone.body.backbone[-1].conv[-1].activation.to("cpu")

# 训练时保持其他部分在MPS运行，仅该层在CPU计算
# 注意：这种方式会增加数据传输开销，影响训练速度

补充说明 #

• 方案1是最快捷的解决方式，无需额外代码修改，适合快速验证训练流程；
• 方案2适合必须使用MobileNetV3 backbone的场景，但需要确保自定义实现的数值精度与原生Hardswish一致；
• 方案3是长期解决方案，升级后能获得MPS后端更完善的算子支持。

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