npu.distribute.npu_distributed_keras_optimizer_wrapper
功能说明
在更新梯度前，添加npu的allreduce操作对梯度进行聚合，然后再更新梯度。该接口仅在分布式场景下使用。
函数原型
def npu_distributed_keras_optimizer_wrapper(optimizer, reduce_reduction="mean", fusion=1, fusion_id=-1, group="hccl_world_group")
参数说明
参数名 |
输入/输出 |
描述 |
|---|---|---|
optimizer |
输入 |
TensorFlow梯度训练优化器。 |
reduction |
输入 |
String类型。 聚合运算的类型，可以为"mean","max","min","prod"或"sum"。 默认为"mean"。 |
fusion |
输入 |
int类型。 allreduce算子融合标识。
|
fusion_id |
输入 |
int类型。 allreduce算子的融合id。 对相同fusion_id的allreduce算子进行融合。 |
group |
输入 |
String类型，最大长度为128字节，含结束符。 group名称，可以为用户自定义group或者"hccl_world_group"。 |
返回值
返回输入的optimizer。
调用示例
1 2 3 4 |
import npu_device as npu optimizer = tf.keras.optimizers.SGD() optimizer = npu.distribute.npu_distributed_keras_optimizer_wrapper(optimizer) # 使用NPU分布式计算，更新梯度 model.compile（optimizer = optimizer） |
