写给新手的 tensorflow:昇腾 TensorFlow 适配到底是啥?

发布时间:2026/5/23 21:19:18

写给新手的 tensorflow:昇腾 TensorFlow 适配到底是啥? 之前做 TensorFlow 模型迁移兄弟问我“哥我们的 TensorFlow 模型能在昇腾上跑吗还是要全部重写”我说能用 tensorflow 适配。好问题。今天一次说清楚。tensorflow 是啥tensorflow 是 TensorFlow 官方的昇腾适配。让你用原生 TensorFlow 接口跑昇腾不用改代码。一句话说清楚tensorflow 是 TensorFlow 官方的昇腾 NPU 适配让你在昇腾上直接用tf.xxx接口不用魔改。你说气人不气人之前要改 TensorFlow 源码才能用昇腾现在一行代码都不用改。为什么要用 tensorflow三个字原生支持。不用 tensorflow魔改版# 之前要用魔改版的 TensorFlowimporttensorflow_npu# 魔改版# 有些接口不兼容modeltensorflow_npu.NPUModel(...)# 特殊接口# 一些功能用不了# tf.distribute.MirroredStrategy() # 不支持# tf.saved_model.save() # 要额外配置用 tensorflow官方版# 现在用官方 TensorFlowimporttensorflowastf# 官方版# 完全原生接口modeltf.keras.Sequential([...])# 标准 TensorFlow# 所有功能都支持strategytf.distribute.MirroredStrategy()tf.saved_model.save(model,model)你说气人不气人现在昇腾和 GPU 的差距就是一个后端。核心概念就三个1. NPU 后端tensorflow 注册了npu后端importtensorflowastf# 检查是否有 npu 后端print(NPU available:,tf.config.list_physical_devices(NPU))# 创建 NPU 张量xtf.constant([1,2,3],dtypetf.float32)xtf.identity(x)# 搬到 NPU自动# 指定设备withtf.device(/NPU:0):ytf.matmul(x,x)print(y)2. 设备映射tensorflow 自动映射设备和内存importtensorflowastf# 设备映射# /NPU:0 → 昇腾 NPU 0 号设备# /GPU:0 → NVIDIA GPU 0 号设备# /CPU:0 → CPU# 自动选择设备devicestf.config.list_physical_devices()print(Available devices:,devices)# 张量设备xtf.constant([1,2,3])print(x.device)# 空在 CPU 上withtf.device(/NPU:0):ytf.constant([1,2,3])print(y.device)# /NPU:0# 模型设备modeltf.keras.Sequential([...])modeltf.keras.models.load_model(model)3. 内存管理tensorflow 自动管理 NPU 内存importtensorflowastf# 自动内存复用# tensorflow 自动# 1. 分配和释放内存# 2. 内存碎片整理# 3. 显存缓存# 手动控制显存缓存gpustf.config.list_physical_devices(NPU)ifgpus:try:# 限制显存使用量tf.config.set_logical_device_configuration(gpus[0],[tf.config.LogicalDeviceConfiguration(memory_limit8192)])# 8GBexceptRuntimeErrorase:print(e)# 查看显存# 用 npu-smi 查看为什么要用 tensorflow三个理由1. 代码不用改原来 GPU 的代码搬到昇腾只要改一个字符串# GPU 代码withtf.device(/GPU:0):modeltf.keras.Sequential([...])# 昇腾代码只改一个字符串withtf.device(/NPU:0):modeltf.keras.Sequential([...])2. 功能全支持TensorFlow 的新功能tensorflow 都支持importtensorflowastf# tf.distribute分布式训练strategytf.distribute.MirroredStrategy()withstrategy.scope():modeltf.keras.Sequential([...])# tf.saved_model模型保存tf.saved_model.save(model,model)# tf.keras.applications预训练模型modeltf.keras.applications.ResNet50(weightsNone,input_shape(224,224,3))# tf.data数据管道datasettf.data.Dataset.from_tensor_slices((x_train,y_train))datasetdataset.batch(32).repeat()3. 性能不差tensorflow 的性能和魔改版差不多importtensorflowastfimporttime# 创建模型modeltf.keras.Sequential([tf.keras.layers.Dense(4096,activationrelu),tf.keras.layers.Dense(4096,activationrelu),tf.keras.layers.Dense(10,activationsoftmax)])model.compile(optimizeradam,losssparse_categorical_crossentropy)# 生成数据xtf.random.normal((1024,4096))ytf.random.uniform((1024,),maxval10,dtypetf.int32)# 预热model.fit(x,x,batch_size32,epochs1,verbose0)# 测性能starttime.time()model.fit(x,x,batch_size32,epochs3,verbose0)elapsedtime.time()-startprint(fTime:{elapsed:.2f}s)print(fThroughput:{1024*3/elapsed:.0f}samples/sec)你说气人不气人TensorFlow 官方支持用起来和 GPU 一样。怎么用代码示例示例 1基础推理importtensorflowastf# 检查 NPU 可用print(NPU available:,tf.config.list_physical_devices(NPU))# 创建模型modeltf.keras.Sequential([tf.keras.layers.Dense(256,activationrelu,input_shape(784,)),tf.keras.layers.Dense(10,activationsoftmax)])# 创建输入xtf.random.normal((32,784))# 推理withtf.device(/NPU:0):outputmodel(x)print(fOutput shape:{output.shape})示例 2训练importtensorflowastf# 数据(x_train,y_train),(x_test,y_test)tf.keras.datasets.mnist.load_data()x_trainx_train.reshape(-1,784).astype(float32)/255.0x_testx_test.reshape(-1,784).astype(float32)/255.0# 模型modeltf.keras.Sequential([tf.keras.layers.Dense(256,activationrelu,input_shape(784,)),tf.keras.layers.Dense(10,activationsoftmax)])model.compile(optimizeradam,losssparse_categorical_crossentropy,metrics[accuracy])# 训练withtf.device(/NPU:0):historymodel.fit(x_train,y_train,batch_size32,epochs3,validation_split0.1)# 评估test_loss,test_accmodel.evaluate(x_test,y_test,batch_size32)print(fTest accuracy:{test_acc:.4f})示例 3分布式训练importosimporttensorflowastf# 环境变量设置os.environ[TF_CONFIG]{cluster: {worker: [localhost:12345]}, task: {type: worker, index: 0}}# 分布式策略strategytf.distribute.MirroredStrategy()withstrategy.scope():# 在策略范围内创建模型modeltf.keras.Sequential([tf.keras.layers.Dense(256,activationrelu,input_shape(784,)),tf.keras.layers.Dense(10,activationsoftmax)])model.compile(optimizeradam,losssparse_categorical_crossentropy,metrics[accuracy])# 数据(x_train,y_train),_tf.keras.datasets.mnist.load_data()x_trainx_train.reshape(-1,784).astype(float32)/255.0# 训练model.fit(x_train,y_train,batch_size32*strategy.num_replicas_in_sync,# 缩放 batch sizeepochs3)示例 4模型保存和加载importtensorflowastf# 创建模型modeltf.keras.Sequential([tf.keras.layers.Dense(256,activationrelu,input_shape(784,)),tf.keras.layers.Dense(10,activationsoftmax)])# 编译model.compile(optimizeradam,losssparse_categorical_crossentropy,metrics[accuracy])# 训练简单训练一下xtf.random.normal((1000,784))ytf.random.uniform((1000,),maxval10,dtypetf.int32)model.fit(x,y,batch_size32,epochs1,verbose0)# 保存模型model.save(/tmp/mnist_model)print(Model saved.)# 加载模型loaded_modeltf.keras.models.load_model(/tmp/mnist_model)print(Model loaded.)# 推理test_inputtf.random.normal((1,784))outputloaded_model(test_input)print(fOutput shape:{output.shape})性能数据在昇腾 910 上对比 GPU操作A100 GPUAscend 910备注推理 (ResNet50)4.5ms4.8ms差不多训练 (batch32)120ms130ms略慢分布式NCCLHCCL都支持模型保存支持支持都支持你说气人不气人现在昇腾和 GPU 差距已经很小了。跟其他仓库的关系tensorflow 在 CANN 架构里属于TensorFlow 官方适配是昇腾对接 TensorFlow 的桥梁。依赖关系TensorFlow官方框架 ↓ 适配 tensorflow昇腾适配 ↓ 调用 hccl / hcomm通信 ↓ 调用 硬件昇腾 NPU解释一下TensorFlow官方深度学习框架tensorflow昇腾适配层hccl / hcomm昇腾通信库硬件昇腾 NPU简单说tensorflow是 TensorFlow 和昇腾之间的桥梁。tensorflow 的核心能力1. 张量操作importtensorflowastf# 创建 NPU 张量xtf.constant([1,2,3],dtypetf.float32)withtf.device(/NPU:0):ytf.matmul(x,x)2. 模型操作importtensorflowastf# 模型迁移到 NPUwithtf.device(/NPU:0):modeltf.keras.Sequential([...])# 模型保存和加载model.save(model)loaded_modeltf.keras.models.load_model(model)3. 分布式importtensorflowastf# 分布式策略strategytf.distribute.MirroredStrategy()withstrategy.scope():modeltf.keras.Sequential([...])4. 数据处理importtensorflowastf# 数据管道datasettf.data.Dataset.from_tensor_slices((x_train,y_train))datasetdataset.batch(32).repeat()# 在 NPU 上训练model.fit(dataset,epochs3)适用场景什么情况下用 tensorflowTensorFlow 迁移GPU 代码迁到昇腾原生支持想要 TensorFlow 官方支持工业部署TensorFlow Serving什么情况下不用PyTorch 项目用 torchtitan-npu极致性能用算子库更底层总结tensorflow 就是 TensorFlow 的昇腾适配原生接口和 GPU 完全一样的接口功能全分布式、模型保存都支持代码不改GPU 代码迁移只要改一个字符串

相关新闻