PyTorch生态与昇腾平台适配：环境搭建与详细安装指南
PyTorch作为业界主流的深度学习框架，广泛应用于深度学习模型的开发与训练。传统上PyTorch默认运行在CPU/GPU环境，而华为昇腾AI处理器通过异构计算架构CANN（Compute Architecture for Neural Networks）为PyTorch提供了强大的算力支持。本文将详细介绍如何在昇腾平台上搭建完整的PyTorch训练环境。底层：昇腾AI处理器+驱动固件中间层：CA
·
文章目录
前言
提示：这里可以添加本文要记录的大概内容：
PyTorch作为业界主流的深度学习框架，广泛应用于深度学习模型的开发与训练。传统上PyTorch默认运行在CPU/GPU环境，而华为昇腾AI处理器通过异构计算架构CANN（Compute Architecture for Neural Networks）为PyTorch提供了强大的算力支持。本文将详细介绍如何在昇腾平台上搭建完整的PyTorch训练环境。
环境架构概述
在昇腾AI处理器上运行PyTorch训练任务，需要构建从底层硬件到上层框架的完整软件栈：
- 底层：昇腾AI处理器+驱动固件
- 中间层：CANN异构计算架构
- 上层：PyTorch框架+适配插件
┌─────────────────────────────────────────┐
│ 应用层 │
│ PyTorch训练脚本 + 自定义模型 │
├─────────────────────────────────────────┤
│ 框架层 │
│ PyTorch + torch_npu适配插件 │
├─────────────────────────────────────────┤
│ CANN层 │
│ Ascend Computing Language │
├─────────────────────────────────────────┤
│ 驱动层 │
│ NPU驱动 + 固件 │
├─────────────────────────────────────────┤
│ 硬件层 │
│ 昇腾AI处理器 │
└─────────────────────────────────────────┘
环境准备与检查
NPU设备检查
在安装驱动和固件前，首先需要确认NPU设备状态及系统兼容性。以Atlas 800训练服务器（型号：9010，Ascend 910处理器）为例：
# 检查NPU是否在位
lspci | grep d801
如果服务器配备N路NPU，将显示N行包含"d801"字段的信息，表明NPU设备正常识别。
0000:89:00.0 Processing accelerators: Device d801
0000:8a:00.0 Processing accelerators: Device d801
0000:8b:00.0 Processing accelerators: Device d801
0000:8c:00.0 Processing accelerators: Device d801
0000:8d:00.0 Processing accelerators: Device d801
0000:8e:00.0 Processing accelerators: Device d801
0000:8f:00.0 Processing accelerators: Device d801
0000:90:00.0 Processing accelerators: Device d801
操作系统兼容性
确保操作系统版本与内核版本符合CANN软件的配套要求，具体可参考昇腾社区文档中的兼容性列表。
驱动与固件安装
创建专用用户
为保障系统安全，建议创建专用用户运行NPU相关服务：
# 创建用户组和用户
groupadd -g 1000 HwHiAiUser
useradd -g HwHiAiUser -u 1000 -d /home/HwHiAiUser -m HwHiAiUser -s /bin/bash
安装驱动
- 下载驱动包：从昇腾社区"固件与驱动"页面下载对应版本的驱动软件包
- 安装准备：
# 添加执行权限
chmod +x Ascend-hdk-910-npu-driver_23.0.rc1_linux-x86-64.run
- 执行安装：
# 完整安装（需要root权限）
./Ascend-hdk-910-npu-driver_23.0.rc1_linux-x86-64.run --full --install-for-all
成功安装后将显示：Driver package installed successfully!
- 验证安装：
npu-smi info
该命令将显示NPU设备状态信息，确认驱动加载正常。
+----------------------------------------------------------------------------------------+
| npu-smi 21.0.1.1.241 Version: 21.0.1.1.241 |
+-------------------+-----------------+-----------------------------------------------------+
| NPU Name | Health | Power(W) Temp(C) Hugepages-Usage(page) |
| Chip | Bus-Id | AICore(%) Memory-Usage(MB) |
+===================+=================+=====================================================+
| 0 910Pro | OK | 48.4 38 0 / 0 |
| 0 | 0000:89:00.0 | 0 0 / 15169 |
+===================+=================+=====================================================+
| 1 910Pro | OK | 48.4 38 0 / 0 |
| 1 | 0000:8A:00.0 | 0 0 / 15169 |
+===================+=================+=====================================================+
安装固件
- 准备固件包：
chmod +x Ascend-hdk-910-npu-firmware_6.3.0.1.241.run
- 执行安装：
./Ascend-hdk-910-npu-firmware_6.3.0.1.241.run --full
成功安装后将显示：Firmware package installed successfully!
系统重启
安装完成后必须重启系统使驱动生效：
reboot
CANN软件依赖安装
系统级依赖
Ubuntu/Debian系列：
apt-get install -y gcc g++ make cmake zlib1g zlib1g-dev openssl libsqlite3-dev \
libssl-dev libffi-dev unzip pciutils net-tools libblas-dev gfortran libblas3
openEuler/CentOS系列：
yum install -y gcc gcc-c++ make cmake unzip zlib-devel libffi-devel openssl-devel \
pciutils net-tools sqlite-devel lapack-devel gcc-gfortran
Python环境配置
以Python 3.7.5为例：
- 下载源码：
wget https://www.python.org/ftp/python/3.7.5/Python-3.7.5.tgz
- 解压与编译：
tar -zxvf Python-3.7.5.tgz
cd Python-3.7.5
./configure --prefix=/usr/local/python3.7.5 --enable-loadable-sqlite-extensions --enable-shared
make -j$(nproc)
make install
- 配置环境变量：
export LD_LIBRARY_PATH=/usr/local/python3.7.5/lib:$LD_LIBRARY_PATH
export PATH=/usr/local/python3.7.5/bin:$PATH
建议将上述环境变量配置添加到~/.bashrc或/etc/profile中实现持久化。
- 验证安装：
python3 --version # 应输出: Python 3.7.5
pip3 --version # 应显示pip版本信息
- 安装Python依赖：
pip3 install attrs numpy decorator sympy cffi pyyaml pathlib2 psutil protobuf scipy requests absl-py
CANN开发套件安装
-
获取安装包：从昇腾社区"CANN"产品页下载对应架构的CANN工具包
-
安装步骤：
# 添加执行权限
chmod +x Ascend-cann-toolkit_6.3.RC1_linux-x86_64.run
# 校验完整性
./Ascend-cann-toolkit_6.3.RC1_linux-x86_64.run --check
# 执行安装
./Ascend-cann-toolkit_6.3.RC1_linux-x86_64.run --install --install-for-all
安装成功后显示：[INFO] xxx install success
[INFO] Ascend-cann-toolkit installed successfully!
[INFO] Please set environment variables by referring to /usr/local/Ascend/ascend-toolkit/set_env.sh
PyTorch环境搭建
前置依赖安装
pip3 install wheel typing_extensions
PyTorch 1.8.1安装指南
x86_64架构：
# 安装官方PyTorch
wget https://download.pytorch.org/whl/cpu/torch-1.8.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
pip3 install torch-1.8.1+cpu-cp37-cp37m-linux_x86_64.whl
# 安装昇腾适配插件
wget https://gitee.com/ascend/pytorch/releases/download/v5.0.rc1-pytorch1.8.1/torch_npu-1.8.1.post1-cp37-cp37m-linux_x86_64.whl
pip3 install torch_npu-1.8.1.post1-cp37-cp37m-linux_x86_64.whl
aarch64架构：
# 安装官方PyTorch
wget https://repo.huaweicloud.com/kunpeng/archive/Ascend/PyTorch/torch-1.8.1-cp37-cp37m-linux_aarch64.whl
pip3 install torch-1.8.1-cp37-cp37m-linux_aarch64.whl
# 安装昇腾适配插件
wget https://gitee.com/ascend/pytorch/releases/download/v5.0.rc1-pytorch1.8.1/torch_npu-1.8.1.post1-cp37-cp37m-linux_aarch64.whl
pip3 install torch_npu-1.8.1.post1-cp37-cp37m-linux_aarch64.whl
安装torchvision：
pip3 install torchvision==0.9.1
PyTorch 1.11.0安装指南
x86_64架构：
# 安装官方PyTorch
wget https://download.pytorch.org/whl/cpu/torch-1.11.0%2Bcpu-cp37-cp37m-linux_x86_64.whl
pip3 install torch-1.11.0+cpu-cp37-cp37m-linux_x86_64.whl
# 安装昇腾适配插件
wget https://gitee.com/ascend/pytorch/releases/download/v5.0.rc1-pytorch1.11.0/torch_npu-1.11.0-cp37-cp37m-linux_x86_64.whl
pip3 install torch_npu-1.11.0-cp37-cp37m-linux_x86_64.whl
aarch64架构：
# 安装官方PyTorch
wget https://repo.huaweicloud.com/kunpeng/archive/Ascend/PyTorch/torch-1.11.0-cp37-cp37m-linux_aarch64.whl
pip3 install torch-1.11.0-cp37-cp37m-linux_aarch64.whl
# 安装昇腾适配插件
wget https://gitee.com/ascend/pytorch/releases/download/v5.0.rc1-pytorch1.11.0/torch_npu-1.11.0-cp37-cp37m-linux_aarch64.whl
pip3 install torch_npu-1.11.0-cp37-cp37m-linux_aarch64.whl
安装torchvision：
pip3 install torchvision==0.12.0
验证安装
执行以下命令验证PyTorch与NPU适配是否成功：
import torch
import torch_npu
# 测试基础功能
a = torch.ones(3, 4).npu()
result = a + a
print(result)
预期输出：
tensor([[2., 2., 2., 2.],
[2., 2., 2., 2.],
[2., 2., 2., 2.]], device='npu:0')
APEX混合精度模块安装
APEX模块为昇腾平台提供混合精度训练支持，可显著提升训练性能。
系统依赖
Ubuntu系统：
apt-get install -y patch build-essential libbz2-dev libreadline-dev wget curl \
llvm libncurses5-dev libncursesw5-dev xz-utils tk-dev liblzma-dev m4 dos2unix \
libopenblas-dev git
CentOS系统：
yum install -y patch libjpeg-turbo-devel dos2unix openblas git
源码编译与安装
- 获取源码：
# 克隆昇腾适配版APEX
git clone -b master https://gitee.com/ascend/apex.git
cd apex
# 获取原生APEX代码
git clone https://github.com/NVIDIA/apex.git
- 切换分支：
cd apex
git checkout 4ef930c1c884fdca5f472ab2ce7cb9b505d26c1a
cd ..
- 生成适配代码：
cd scripts
bash gen.sh
- 编译安装包：
cd ../apex
python3 setup.py --cpp_ext --npu_float_status bdist_wheel
- 安装APEX：
x86_64架构：
cd dist
pip3 install apex-0.1_ascend-cp37-cp37m-linux_x86_64.whl
aarch64架构：
cd dist
pip3 install apex-0.1_ascend-cp37-cp37m-linux_aarch64.whl
完整环境验证
创建测试脚本test_npu_environment.py：
import torch
import torch_npu
from apex import amp
def test_environment():
print("=== PyTorch与昇腾NPU环境验证 ===")
# 检查PyTorch版本
print(f"PyTorch版本: {torch.__version__}")
# 检查NPU设备
if torch.npu.is_available():
print(f"NPU设备数量: {torch.npu.device_count()}")
for i in range(torch.npu.device_count()):
print(f"NPU设备 {i}: {torch.npu.get_device_name(i)}")
else:
print("NPU设备不可用")
return False
# 基础张量运算测试
print("\n=== 基础运算测试 ===")
x = torch.randn(2, 3).npu()
y = torch.randn(2, 3).npu()
z = x + y
print(f"张量加法测试: {z.shape}, device: {z.device}")
# 模型前向传播测试
print("\n=== 模型推理测试 ===")
model = torch.nn.Linear(10, 5).npu()
input_data = torch.randn(1, 10).npu()
output = model(input_data)
print(f"模型推理输出: {output.shape}")
print("\n✅ 环境验证通过！PyTorch与昇腾NPU适配成功")
return True
if __name__ == "__main__":
test_environment()
运行结果示例：
=== PyTorch与昇腾NPU环境验证 ===
PyTorch版本: 1.8.1
NPU设备数量: 8
NPU设备 0: Ascend 910
=== 基础运算测试 ===
张量加法测试: torch.Size([2, 3]), device: npu:0
=== 模型推理测试 ===
模型推理输出: torch.Size([1, 5])
✅ 环境验证通过！PyTorch与昇腾NPU适配成功
结语
现在已完成昇腾平台上完整的PyTorch训练环境搭建。这个环境支持：
- 原生PyTorch API与编程范式
- NPU设备加速计算
- 混合精度训练优化
- 分布式训练支持
开发者现在可以将现有的PyTorch训练脚本迁移到昇腾平台，只需将设备指定为NPU即可享受昇腾处理器的强大算力：
# 将cuda设备切换为npu
device = torch.device("npu:0")
model = model.to(device)
data = data.to(device)
# 或者直接使用npu()方法
model = model.npu()
data = data.npu()
性能优化建议
- 使用混合精度训练：
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1", loss_scale=128.0)
- 启用数据并行：
if torch.npu.device_count() > 1:
model = torch.nn.DataParallel(model)
资源与支持
昇腾RAG SDK技术原理和功能介绍
慕课：https://www.hiascend.com/developer/courses/detail/1980194995581026305
Pytorch小模型迁移与精度调优
慕课：https://www.hiascend.com/developer/courses/detail/1983449864584704001
有“AI”的1024 = 2048，欢迎大家加入2048 AI社区
更多推荐
18
0- 0
已为社区贡献14条内容
所有评论(0)