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

相关推荐
牛奶几秒前
AI辅助开发的基础概念
前端·人工智能·ai编程
东坡肘子1 分钟前
OpenClaw 不错,但我好像没有那么需要 -- 肘子的 Swift 周报 #125
人工智能·swiftui·swift
风象南8 小时前
普通人用AI加持赚到的第一个100块
人工智能·后端
牛奶8 小时前
2026年大模型怎么选?前端人实用对比
前端·人工智能·ai编程
牛奶8 小时前
前端人为什么要学AI?
前端·人工智能·ai编程
罗西的思考11 小时前
AI Agent框架探秘:拆解 OpenHands(10)--- Runtime
人工智能·算法·机器学习
冬奇Lab12 小时前
OpenClaw 源码精读(2):Channel & Routing——一条消息如何找到它的 Agent?
人工智能·开源·源码阅读
冬奇Lab12 小时前
一天一个开源项目(第38篇):Claude Code Telegram - 用 Telegram 远程用 Claude Code,随时随地聊项目
人工智能·开源·资讯
孟健13 小时前
Karpathy 用 200 行纯 Python 从零实现 GPT:代码逐行解析
python
格砸13 小时前
从入门到辞职|从ChatGPT到OpenClaw,跟上智能时代的进化
前端·人工智能·后端