KerasDistributeOptimizer构造函数

功能说明

KerasDistributeOptimizer类的构造函数，用于包装用户使用tf.Keras构造的脚本中的单机训练优化器，构造NPU分布式训练优化器。

函数原型

def __init__(self, optimizer, name="NpuKerasOptimizer", **kwargs)

参数说明

返回值

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