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())) |