---
title: npu_distributed_optimizer_wrapper
description: "| 产品 | 是否支持 |"
url: https://www.hiascend.com/document/detail/zh/TensorFlowCommercial/latest/migration/tfmigr1/tfmigr1_tfadapi_0059.html
sourcePath: /source/zh/TensorFlowCommercial/900/migration/tfmigr1/tfmigr1_tfadapi_0059.html
indexId: ec1aae2a06aa72857288d6fa50140f90e5d9679bac678b5377761a179293bafa79
---
# npu_distributed_optimizer_wrapper

#### 产品支持情况

| 产品 | 是否支持 |
| --- | --- |
| Atlas 350 加速卡 | √ |
| Atlas A3 训练系列产品 / Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品 / Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | ☓ |
| Atlas 推理系列产品 | ☓ |
| Atlas 训练系列产品 | √ |


#### 功能说明

对传入的optimizer中的求梯度的函数添加NPU的allreduce操作之后，将包含原生优化器求梯度和NPU的allreduce两个操作合并为一个函数，替换原生优化器的求梯度的函数，最终返回输入的优化器。该接口仅在分布式场景下使用。


#### 函数原型

```
def npu_distributed_optimizer_wrapper(optimizer)
```


#### 参数说明

| 参数名 | 输入/输出 | 描述 |
| --- | --- | --- |
| optimizer | 输入 | TensorFlow梯度训练优化器。 |


#### 返回值

返回输入的optimizer。


#### 调用示例

```
from npu_bridge.npu_init import *
# tf场景
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) # 使用SGD优化器
optimizer = npu_distributed_optimizer_wrapper(optimizer) # 使用NPU分布式计算，更新梯度
# keras场景
optimizer = tf.keras.optimizers.SGD()
optimizer = npu_distributed_optimizer_wrapper(optimizer) # 使用NPU分布式计算，更新梯度
```