

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

# 在 TensorFlow 訓練指令碼中使用 SMDDP 程式庫 (已棄用)
<a name="data-parallel-modify-sdp-tf2"></a>

**重要**  
SMDDP 程式庫已停止對 TensorFlow 的支援，且不再於 v2.11.0 之後的 TensorFlow DLC 中提供。若要尋找已安裝 SMDDP 程式庫的舊版 TensorFlow DLC，請參閱[支援的架構](distributed-data-parallel-support.md#distributed-data-parallel-supported-frameworks)。

下列步驟說明如何修改 TensorFlow 訓練指令碼，利用 SageMaker AI 的分散式資料平行化資料庫。  

程式庫 API 的設計旨在近似 Horovod API。請參閱 [SageMaker AI 分散式資料平行化 TensorFlow API 文件](https://sagemaker.readthedocs.io/en/stable/api/training/smd_data_parallel.html#api-documentation)，了解程式庫為 TensorFlow 提供之每個 API 的其他詳細資訊。

**注意**  
SageMaker AI 分散式資料平行化可配合包含 `tf` 核心模組的 TensorFlow 訓練指令碼 (`tf.keras` 模組除外) 調整。SageMaker AI 分散式資料平行化不支援包含 Keras 實作的 TensorFlow。

**注意**  
SageMaker AI 分散式資料平行化程式庫支援開箱即用的自動混合精確度 (AMP)。除了對訓練指令碼進行架構層級修改之外，無需執行任何額外操作即可啟用 AMP。如果梯度是 FP16，SageMaker AI 資料平行化程式庫會以 FP16 執行其 `AllReduce` 操作。有關將 AMP API 實施到訓練腳本的詳細資訊，請參閱以下資源：  
[架構 - ](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#tensorflow)*NVIDIA 深度學習效能文件中的 TensorFlow*
*NVIDIA 開發人員文件*中[適用於深度學習的自動混合精確度](https://developer.nvidia.com/automatic-mixed-precision)
*TensorFlow 文件*中的 [TensorFlow 混合精確度 API](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，工作者節點是等級 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) 到所有工作者節點 (排名 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. 最後，修改指令碼，僅在領導節點保存檢查點。領導節點有同步化的模型。這也可避免工作者節點覆寫檢查點，以及可能損壞檢查點。

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

完成訓練指令碼的調整後，請繼續前往 [使用 SageMaker Python SDK 透過 SMDDP 啟動分散式訓練任務](data-parallel-use-api.md)。