npu.distribute.npu_distributed_keras_optimizer_wrapper

函数原型

def npu_distributed_keras_optimizer_wrapper(optimizer, reduce_reduction="mean", fusion=1, fusion_id=-1, group="hccl_world_group")

功能说明

在更新梯度前，添加npu的allreduce操作对梯度进行聚合，然后再更新梯度。该接口仅在分布式场景下使用。

参数说明

参数名	输入/输出	描述
optimizer	输入	TensorFlow梯度训练优化器。
reduction	输入	String类型。聚合运算的类型，可以为"mean","max","min","prod"或"sum"。默认为"mean"。
fusion	输入	int类型。 allreduce算子融合标识。 0：不融合，该allreduce算子不和其他allreduce算子融合。 1：按照梯度切分策略进行融合，默认为1。 2：按照相同fusion_id进行融合。
fusion_id	输入	int类型。 allreduce算子的融合id。对相同fusion_id的allreduce算子进行融合。
group	输入	String类型，最大长度为128字节，含结束符。 group名称，可以为用户自定义group或者"hccl_world_group"。

返回值

返回输入的optimizer。

调用示例

import npu_device as npu
optimizer = tf.keras.optimizers.SGD()
optimizer = npu_distributed_keras_optimizer_wrapper(optimizer) # 使用NPU分布式计算，更新梯度
model.compile（optimizer = optimizer）

父主题： TF Adapter 2.6 接口参考