如何在 Npu 环境中安装 mmcv

Source: Dev.to

描述

我想在 Ascend GPU 上运行 musetalk 模型，但 mmcv 带来了很多麻烦。我总结了一些经验，想在这里分享。

我的环境

• NPU: 910b2
• OS: aarch64 Ubuntu 22.04.5 LTS

故事

当我尝试使用 pip 安装 mmcv 时，出现了 No module named mmcv._ext 错误。只能通过从源码编译来解决。编译在 CPU 环境下成功，但在 NPU 环境下失败。

在 CPU 上编译

从 GitHub 下载 mmcv 代码，分支 npu-dev：

https://github.com/open-mmlab/mmcv/tree/npu-dev

修改 setup.py

第 21 行

# cmd_class = {'build_ext': BuildExtension}
cmd_class = {'build_ext': BuildExtension.with_options(use_ninja=False)}

第 273‑300 行

elif (os.getenv('FORCE_NPU', '0') == '1'):
    print(f'Compiling {ext_name} only with CPU and NPU')
    try:
        from torch_npu.utils.cpp_extension import NpuExtension
        define_macros += [('MMCV_WITH_NPU', None)]
        extension = NpuExtension
    except Exception:
        raise ImportError('can not find any torch_npu')
    # src
    op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
        glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
        glob.glob('./mmcv/ops/csrc/common/npu/*.cpp') + \
        glob.glob('./mmcv/ops/csrc/pytorch/npu/*.cpp')
    include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
    include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/npu'))
    if 'cxx' not in extra_compile_args:
        extra_compile_args['cxx'] = {}
    extra_compile_args['cxx'] += ['-std=c++17', '-Wall']

else:
    print(f'Compiling {ext_name} only with CPU')
    op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
        glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp')
    extension = CppExtension
    include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
    if 'cxx' not in extra_compile_args:
        extra_compile_args['cxx'] = {}
    extra_compile_args['cxx'] += ['-std=c++17', '-Wall']

构建命令

前置条件： 在运行以下命令前请先安装 torch 和 torchvision。

MMCV_WITH_OPS=1 MAX_JOBS=40 python setup.py build_ext

如果运气好，你也可以尝试：

MMCV_WITH_OPS=1 MAX_JOBS=40 python setup.py develop

注意： develop 会在 site‑packages 文件夹创建符号链接。

验证 mmcv 模块：

from mmcv.ops import MultiScaleDeformableAttention

实际的安装命令：

MMCV_WITH_OPS=1 MAX_JOBS=40 python setup.py install

注意

文件 _ext.cpython-310-aarch64-linux-gnu.so 可能不在 mmcv 目录下，而是位于构建目录中。你必须手动创建符号链接。