KerasDistributeOptimizer构造函数
函数原型
def __init__(self, optimizer, name="NpuKerasOptimizer", **kwargs)
功能说明
KerasDistributeOptimizer类的构造函数,用于包装用户使用tf.Keras构造的脚本中的单机训练优化器,构造NPU分布式训练优化器。
参数说明
参数名 |
输入/输出 |
描述 |
---|---|---|
optimizer |
输入 |
用于梯度计算和更新权重的单机版训练优化器。 |
name |
输入 |
优化器名称。 |
返回值
返回KerasDistributeOptimizer类对象。
调用示例
1 2 3 4 5 |
import tensorflow as tf from npu_bridge.npu_init import * model=xxx model.compile(loss='mean_squared_error', optimizer=KerasDistributeOptimizer(tf.keras.optimizers.SGD())) |