

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

# 在 TensorFlow 训练脚本中使用 SMDDP 库（已弃用）
<a name="data-parallel-modify-sdp-tf2"></a>

**重要**  
在 v2.11.0 之后，SMDDP 库已停止支持， TensorFlow 并且在 DLCs v2.11.0 TensorFlow 之后不再可用。要查找以前安装了 TensorFlow DLCs SMDDP 库的情况，请参阅。[支持的框架](distributed-data-parallel-support.md#distributed-data-parallel-supported-frameworks)

以下步骤向您展示如何修改 TensorFlow 训练脚本以利用 SageMaker AI 的分布式数据 parallel 库。  

该库 APIs 的设计类似于 Horovod APIs。有关该库提供的每个 API 的更多详细信息 TensorFlow，请参阅 [SageMaker AI 分布式数据 parallel TensorFlow API 文档](https://sagemaker.readthedocs.io/en/stable/api/training/smd_data_parallel.html#api-documentation)。

**注意**  
SageMaker AI 分布式数据 parallel 适用于由除`tf.keras`模块之外的`tf`核心模块组成的 TensorFlow 训练脚本。 SageMaker 人工智能分布式数据 parallel 不支持 Ker TensorFlow as 实现。

**注意**  
 SageMaker AI 分布式数据并行度库支持开箱即用的自动混合精度 (AMP)。除了对训练脚本进行框架级别的修改之外，您无需执行额外操作即可启用 AMP。如果有梯度 FP16，则 SageMaker AI 数据并行度库将在中运行其操作。`AllReduce` FP16有关在训练脚本中实现 AMP APIs 的更多信息，请参阅以下资源：  
[框架- TensorFlow](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#tensorflow) 在 *NVIDIA 深度学习性能文档*中
*NVIDIA 开发人员文档*中的[适用于深度学习的自动混合精度](https://developer.nvidia.com/automatic-mixed-precision)
TensorFlow *TensorFlow文档 APIs*中的@@ [精度好坏参半](https://www.tensorflow.org/guide/mixed_precision)

1. 导入库的 TensorFlow 客户端并对其进行初始化。

   ```
   import smdistributed.dataparallel.tensorflow as sdp 
   sdp.init()
   ```

1. 使用 `local_rank` 将每个 GPU 固定到一个 `smdistributed.dataparallel` 进程，这表示进程在给定节点中的相对秩。`sdp.tensorflow.local_rank()` API 向您提供设备的局部秩。领导节点的秩为 0，Worker 节点的秩为 1、2、3，依此类推。在以下代码块中将其调用为`sdp.local_rank()`。 `set_memory_growth`与分布式 SageMaker AI 没有直接关系，但必须将其设置为使用进行分布式训练 TensorFlow。

   ```
   gpus = tf.config.experimental.list_physical_devices('GPU')
   for gpu in gpus:
       tf.config.experimental.set_memory_growth(gpu, True)
   if gpus:
       tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')
   ```

1. 根据工作线程数扩展学习率。`sdp.tensorflow.size()` API 为您提供集群中工作线程的数量。以下代码块中将其作为 `sdp.size()` 调用。

   ```
   learning_rate = learning_rate * sdp.size()
   ```

1. 在训练期间，使用库的 `DistributedGradientTape` 来优化 `AllReduce` 操作。这会包装 `tf.GradientTape`。  

   ```
   with tf.GradientTape() as tape:
         output = model(input)
         loss_value = loss(label, output)
       
   # SageMaker AI data parallel: Wrap tf.GradientTape with the library's DistributedGradientTape
   tape = sdp.DistributedGradientTape(tape)
   ```

1. 将初始模型变量从领导节点（秩 0）广播到所有 Worker 节点（秩 1 到 n）。这是确保对所有工作线程秩进行一致的初始化所必需的。在模型和优化器变量初始化后使用 `sdp.tensorflow.broadcast_variables` API。以下代码块中将其作为 `sdp.broadcast_variables()` 调用。

   ```
   sdp.broadcast_variables(model.variables, root_rank=0)
   sdp.broadcast_variables(opt.variables(), root_rank=0)
   ```

1. 最后，修改脚本，仅在领导节点上保存检查点。领导节点具有同步模型。这还可以避免 Worker 节点覆盖检查点并可能损坏检查点。

   ```
   if sdp.rank() == 0:
       checkpoint.save(checkpoint_dir)
   ```

以下是使用库进行分布式 TensorFlow 训练的示例训练脚本。

```
import tensorflow as tf

# SageMaker AI data parallel: Import the library TF API
import smdistributed.dataparallel.tensorflow as sdp

# SageMaker AI data parallel: Initialize the library
sdp.init()

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    # SageMaker AI data parallel: Pin GPUs to a single library process
    tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')

# Prepare Dataset
dataset = tf.data.Dataset.from_tensor_slices(...)

# Define Model
mnist_model = tf.keras.Sequential(...)
loss = tf.losses.SparseCategoricalCrossentropy()

# SageMaker AI data parallel: Scale Learning Rate
# LR for 8 node run : 0.000125
# LR for single node run : 0.001
opt = tf.optimizers.Adam(0.000125 * sdp.size())

@tf.function
def training_step(images, labels, first_batch):
    with tf.GradientTape() as tape:
        probs = mnist_model(images, training=True)
        loss_value = loss(labels, probs)

    # SageMaker AI data parallel: Wrap tf.GradientTape with the library's DistributedGradientTape
    tape = sdp.DistributedGradientTape(tape)

    grads = tape.gradient(loss_value, mnist_model.trainable_variables)
    opt.apply_gradients(zip(grads, mnist_model.trainable_variables))

    if first_batch:
       # SageMaker AI data parallel: Broadcast model and optimizer variables
       sdp.broadcast_variables(mnist_model.variables, root_rank=0)
       sdp.broadcast_variables(opt.variables(), root_rank=0)

    return loss_value

...

# SageMaker AI data parallel: Save checkpoints only from master node.
if sdp.rank() == 0:
    checkpoint.save(checkpoint_dir)
```

调整完训练脚本后，请继续到[使用 Python SageMaker SDK 使用 SMDDP 启动分布式训练作业](data-parallel-use-api.md)。