npu_distributed_optimizer_wrapper

Applicability

Product	Supported (√/x)
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	X
Atlas inference product	X
Atlas training product	√

Description

Adds the AllReduce operation of NPU to the input gradient function of the optimizer, combines them into one function, and returns the optimizer. This API is used only in distributed scenarios.

Prototype

def npu_distributed_optimizer_wrapper(optimizer)

Parameters

Parameter	Input/Output	Description
optimizer	Input	TensorFlow gradient training optimizer.

Returns

TensorFlow gradient training optimizer.

Example

from npu_bridge.npu_init import *
# TF scenario
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) # Use the SGD optimizer.
optimizer = npu_distributed_optimizer_wrapper(optimizer) # Use NPU-based distributed computing to update gradients.
# Keras scenario
optimizer = tf.keras.optimizers.SGD()
optimizer = npu_distributed_optimizer_wrapper(optimizer) # Use NPU-based distributed computing to update gradients.

Parent topic: npu_bridge.estimator.npu.npu_optimizer