若想将gpu的代码在昇腾npu上运行,创建docker应该创建怎么样的docker?(待完善)
直接把 GPU 的 PyTorch/TensorFlow 代码搬到昇腾 NPU 上运行,通常不能直接复用 NVIDIA CUDA 的 Docker 镜像,而是要用 华为官方提供的 Ascend NPU 适配镜像,里面预装了 CANN、驱动适配层和对应的 PyTorch-NPU/TensorFlow-NPU。
1. 基础原则
GPU Docker: 基于 NVIDIA 的 CUDA 镜像,例如
nvidia/cuda:11.8-cudnn8-devel-ubuntu20.04
,靠 CUDA Toolkit + cuDNN 提供算力支持。NPU Docker (昇腾): 必须基于 华为官方 Ascend 镜像,里面包含 CANN(Compute Architecture for Neural Networks)、NPU 驱动接口,以及对应的 AI 框架(如
torch-npu
、mindspore
)。
如果用 GPU 镜像直接跑,NPU 根本无法识别。
2. 官方推荐的 Docker 镜像获取方式
华为 Ascend Hub
华为在 Ascend Hub 和 CANN 下载中心 提供了官方镜像,比如:ascend-pytorch:5.0.0-py3.8-ubuntu18.04
ascend-mindspore:2.3.1-ubuntu20.04
ascend-tensorflow:2.11.0-ubuntu20.04
这些镜像里已经集成:
CANN Toolkit
对应版本的 PyTorch-NPU / TensorFlow-NPU
驱动交互层
Dockerfile 自建镜像(更灵活)
如果要和现有 GPU 代码环境更接近,可以自己写 Dockerfile:FROM ubuntu:20.04# 基础依赖 RUN apt-get update && apt-get install -y \wget git vim python3-dev python3-pip libglib2.0-dev# 安装 Ascend CANN Toolkit COPY Ascend-cann-toolkit*.run /tmp/ RUN bash /tmp/Ascend-cann-toolkit*.run --install# 配置环境变量 ENV ASCEND_TOOLKIT_HOME=/usr/local/Ascend/ascend-toolkit/latest ENV PATH=$ASCEND_TOOLKIT_HOME/bin:$PATH ENV LD_LIBRARY_PATH=$ASCEND_TOOLKIT_HOME/lib64:$LD_LIBRARY_PATH ENV PYTHONPATH=$ASCEND_TOOLKIT_HOME/python/site-packages:$PYTHONPATH# 安装 PyTorch-NPU RUN pip install torch==2.1.0 RUN pip install torch-npu==2.1.0.post3
这样就能构建一个和 GPU 环境结构类似的 NPU Docker。
3. 如何选择镜像
如果你主要跑 PyTorch 项目(原来 GPU 用
torch
):
→ 用ascend-pytorch
系列镜像,或者自己做基于 Ubuntu + torch-npu 的定制镜像。如果你主要跑 TensorFlow 项目:
→ 用ascend-tensorflow
系列镜像。如果要和 GPU 的训练脚本几乎无缝迁移:
尽量选择 相同 Python 版本、相同框架版本的 ascend 镜像,只需要在代码里做
cuda → npu
的改动。
4. 注意事项
代码层修改:GPU 上
device = torch.device("cuda:0")
,要改为device = torch.device("npu:0")
。依赖库:一些 CUDA 专用库(如 apex、cupy)在 NPU 上不可用,需要替换/删除。
性能优化:昇腾 NPU 对算子支持范围和优化点和 GPU 不一样,可能需要用到 torch_npu.ops 或者 CANN 自定义算子。