基于昇腾 配置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 并正常打印

相关推荐
OceanBase数据库官方博客5 分钟前
OceanBase DataPilot AIP:Ontology 承载AI能力面的另一条路
人工智能·oceanbase
笨鸟先飞,勤能补拙1 小时前
AI 赋能网络安全:技术全景、成熟度评估与实战案例
人工智能·python·安全·web安全·网络安全·sqlite·github
一次旅行1 小时前
AI 前沿日报 | 2026年07月31日
人工智能
2601_963749101 小时前
标题:越华环保集团|面向美丽河湖项目的数字化污水治理云边协同采集架构设计
人工智能
沐籽李1 小时前
从溶剂可及表面积SASA理解抗体结构与工程改造
人工智能·药物设计·aidd·sasa
智慧物业老杨1 小时前
物业如何做好预算管理?落地架构逻辑
人工智能·架构
微学AI2 小时前
一根针指向所有方向:挂谷猜想对 LLM Agent 技能-记忆架构的启示
开发语言·人工智能·架构·挂谷猜想
城管不管2 小时前
ReAct、Plan-and-Execute、Reflection 三大智能 Agent 范式核心区别
java·人工智能·算法·spring·ai·动态规划
boppu2 小时前
布草特殊污渍去渍剂的种类及作用
大数据·人工智能
长和信泰光伏储能2 小时前
京津冀光伏发电:绿色能源的未来之路
python·能源