基于昇腾 配置pytorch环境

目录

一、基础环境

二、基础准备

三、安装pytorch

四、测试


一、基础环境

安装conda

安装配套版本的NPU驱动固件、CANN软件(Toolkit、Kernels和NNAL)并配置CANN环境变量,具体请参见《CANN 软件安装指南》(商用版)或《CANN 软件安装指南》(社区版)。

基础说明:

昇腾软件介绍

软件类型 软件介绍
昇腾NPU固件 固件包含昇腾AI处理器自带的OS 、电源器件和功耗管理器件控制软件,分别用于后续加载到AI处理器的模型计算、芯片启动控制和功耗控制。
昇腾NPU驱动 部署在昇腾服务器,管理查询昇腾AI处理器,同时为上层CANN软件提供芯片控制、资源分配等接口。
CANN 部署在昇腾服务器,包含Runtime、算子库、图引擎、媒体数据处理等组件,通过AscendCL(Ascend Computing Language,昇腾计算语言)对外提供Device管理、Context管理、Stream管理、内存管理、模型加载与执行、算子加载与执行、媒体数据处理等API,帮助开发者实现在昇腾软硬件平台上开发和运行AI业务。 CANN软件按照功能主要分为Toolkit(开发套件)、Kernels(算子包)、NNAE(深度学习引擎)、NNRT(离线推理引擎)、NNAL(加速库)、TFPlugin(TensorFlow框架插件)几种软件包,各软件包支持功能范围如下: * Toolkit:主要用于训练和推理业务、模型转换、算子/应用/模型的开发和编译。 * NNAE:主要用于训练和推理业务。 * NNRT:仅支持离线推理,主要用于应用程序的模型推理。 * TFPlugin:用于运行训练业务时和TensorFlow框架进行对接,帮助TensorFlow框架调用底层CANN接口运行训练业务。 * Kernels:依赖Toolkit、NNRT或NNAE软件包,可以节省算子编译时间。在大模型推理、训练场景,使用静态库编译应用程序,和运行包含动态shape网络或单算子API(例如aclnn类API)相关业务时安装二进制算子包。 * NNAL:依赖Toolkit或NNAE,包含面向大模型领域的ATB(Ascend Transformer Boost)加速库,可以提升大模型训练和推理性能。

二、基础准备

复制代码
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

以上命令以root用户安装后的默认路径为例,请用户根据set_env.sh的实际路径进行替换。

三、安装pytorch

配置conda环境

复制代码
conda create -n pytorch python=3.9
conda activate pytorch 

在昇腾社区选择对应的pytorch版本与torch_npu版本

https://www.hiascend.com/document/detail/zh/Pytorch/60RC3/configandinstg/instg/insg_0005.html

安装pytorch

复制代码
# 下载插件包
wget https://gitcode.com/Ascend/pytorch/releases/download/v7.2.0-pytorch2.1.0/torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
# 安装命令
pip3 install torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

安装torch_npu

复制代码
# 下载插件包
wget https://gitee.com/ascend/pytorch/releases/download/v6.0.rc3-pytorch2.1.0/torch_npu-2.1.0.post8-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
# 安装命令
pip3 install torch_npu-2.1.0.post8-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

四、测试

复制代码
python3 -c "
import torch
import torch_npu
x = torch.randn(2, 3).npu()
y = x + x
print('Success:', y.cpu())
"

结果如下:

✅ x = torch.randn(2, 3).npu() 成功将张量放到 NPU

✅ y = x + x 在 NPU 上完成计算

✅ .cpu() 将结果传回 CPU 并正常打印

相关推荐
武子康25 分钟前
单卡 A6000 跑 Qwen-Drive:怎样留下可比较的第一条基线?
linux·人工智能·深度学习
严肃小白兔25 分钟前
第 06/24 期:神经网络的反向传播算法:从链式法则到梯度计算
人工智能·神经网络·算法
用户0190275816143 分钟前
天勤 tqsdk 能做 A 股行情吗?期货强项与股票数据选型
python
后端小肥肠44 分钟前
公众号漫画又挖到长尾赛道:家常菜美食,流量稳定长期有收益
人工智能·aigc·agent
泯泷1 小时前
从"检索管道"到"主动助理"——Agentic RAG 与双层记忆(五)
人工智能·llm·agent
武子康1 小时前
ComfyUI 加载了动作 LoRA,H3-World 的时间绑定就接好了吗?
人工智能·机器学习·claude
腾科IT教育1 小时前
HCIE-AI认证全部解析:方向选择、考试门槛与职业适配指南
人工智能·华为认证·hcie-ai·华为ai
QXWZ_IA1 小时前
具身智能在油气行业规模化应用面临哪些技术挑战?
人工智能·安全
2601_967979122 小时前
pip在Python中的9个必须知道的用法
python·pip·版本控制·包管理工具·requirements.txt