文档
注册
评分
提单
论坛
小AI

npu_distributed_optimizer_wrapper

函数原型

def npu_distributed_optimizer_wrapper(optimizer)

功能说明

对传入的optimizer中的求梯度的函数添加npu的allreduce操作之后,将包含原生优化器求梯度和npu allreduce两个操作合并为一个函数,替换原生优化器的求梯度的函数,最终返回输入的优化器。该接口仅在分布式场景下使用。

参数说明

参数名

输入/输出

描述

optimizer

输入

TensorFlow梯度训练优化器。

返回值

返回输入的optimizer。

调用示例

from npu_bridge.npu_init import *
# tf场景
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) # 使用SGD优化器
optimizer = npu_distributed_optimizer_wrapper(optimizer) # 使用NPU分布式计算,更新梯度
# keras场景
optimizer = tf.keras.optimizers.SGD()
optimizer = npu_distributed_optimizer_wrapper(optimizer) # 使用NPU分布式计算,更新梯度
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词