

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

# 使用自有深度學習模型
<a name="training-compiler-modify-scripts"></a>

**重要**  
Amazon Web Services (AWS) 宣佈不再推出新版本的 SageMaker Training Compiler。您可以透過現有的 SageMaker 訓練 AWS 深度學習容器 (DLC)，繼續利用 SageMaker Training Compiler。請務必注意，雖然現有 DLCs仍可存取，但根據深度學習容器架構支援政策 AWS，他們將不再收到來自 的修補程式或更新。 [AWS](https://docs.aws.amazon.com/deep-learning-containers/latest/devguide/support-policy.html)

本指南將逐步引導您如何針對利用編譯器加速的訓練工作調整訓練指令碼。訓練指令碼的準備取決於以下內容：
+ 訓練設定，例如單核心或分散式訓練。
+ 用來建立訓練指令碼的架構與程式庫。

根據採用的架構，選擇下列其中一個主題。

**Topics**
+ [PyTorch](training-compiler-pytorch-models.md)
+ [TensorFlow](training-compiler-tensorflow.md)

**注意**  
在準備好訓練指令碼之後，您可以利用 SageMaker AI 架構估算器類別來執行 SageMaker 訓練任務。如需詳細資訊，請前往 [啟用 SageMaker Training Compiler](training-compiler-enable.md) 參閱前一主題。

# PyTorch
<a name="training-compiler-pytorch-models"></a>

攜帶您自己的 PyTorch 模型至 SageMaker AI，並使用 SageMaker Training Compiler 執行訓練任務。

**Topics**
+ [配置 Hugging Face 轉換器的 PyTorch 模型](#training-compiler-pytorch-models-transformers)

## 配置 Hugging Face 轉換器的 PyTorch 模型
<a name="training-compiler-pytorch-models-transformers"></a>

配置 [Hugging Face 轉換器](https://huggingface.co/docs/transformers/index)的 PyTorch 模型是基於 PyTorch 的 [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) API。Hugging Face 轉換器也為 PyTorch 提供[訓練器](https://huggingface.co/docs/transformers/main_classes/trainer)和預先訓練的模型類別，以妥善減少設定自然語言處理 (NLP) 模型的工作量。準備訓練指令碼之後，您可以使用 SageMaker AI `PyTorch` 或 `HuggingFace` 估算器與 SageMaker Training Compiler 組態啟動訓練任務，藉此如 [啟用 SageMaker Training Compiler](training-compiler-enable.md) 所述繼續進行下一個主題的操作。

**提示**  
當您在訓練指令碼中使用轉換器為 NLP 模型建立字符化工具時，請確保您透過指定 `padding='max_length'` 來使用靜態輸入張量形狀。請勿使用 `padding='longest'`，因為填補至批次中最長的序列可能會變更每個訓練批次的張量形狀。動態輸入形狀可觸發模型的重新編譯，並可能會增加總訓練時間。有關轉換器權杖化工具填補選項的詳情，請參閱 *Hugging Face 轉換器文件*中的[填補和截斷](https://huggingface.co/docs/transformers/pad_truncation)。

**Topics**
+ [使用 Hugging Face 轉換器 `Trainer` 類別的大型語言模型](#training-compiler-pytorch-models-transformers-trainer)
+ [直接使用 PyTorch 的大型語言模型 (無需 Hugging Face 轉換器訓練器 API)](#training-compiler-pytorch-models-non-trainer)

### 使用 Hugging Face 轉換器 `Trainer` 類別的大型語言模型
<a name="training-compiler-pytorch-models-transformers-trainer"></a>

如果您使用轉換器程式庫的訓練器類別，則不需對訓練指令碼進行任何其他變更。如果您透過估算器類別啟用訓練器模型，SageMaker Training Compiler 會自動編譯訓練器模型。下列程式碼顯示一個 PyTorch 訓練指令碼與 Hugging Face 訓練器 API 的基本形式。

```
from transformers import Trainer, TrainingArguments

training_args=TrainingArguments(**kwargs)
trainer=Trainer(args=training_args, **kwargs)
```

**Topics**
+ [適用於單一 GPU 訓練](#training-compiler-pytorch-models-transformers-trainer-single-gpu)
+ [適用於分散式訓練](#training-compiler-pytorch-models-transformers-trainer-distributed)
+ [將 SageMaker Training Compiler 搭配使用 `Trainer` 的最佳實務](#training-compiler-pytorch-models-transformers-trainer-best-practices)

#### 適用於單一 GPU 訓練
<a name="training-compiler-pytorch-models-transformers-trainer-single-gpu"></a>

當您使用 [https://huggingface.co/docs/transformers/main_classes/trainer](https://huggingface.co/docs/transformers/main_classes/trainer) 類別時，不需變更程式碼。

#### 適用於分散式訓練
<a name="training-compiler-pytorch-models-transformers-trainer-distributed"></a>

**PyTorch v1.11.0 及較新版本**

若要使用 SageMaker Training Compiler 執行分散式訓練，您必須在訓練指令碼中新增下列 `_mp_fn()` 函式並包裝 `main()` 函式。它會將 `_mp_fn(index)` 函式呼叫從 PyTorch (`pytorchxla`) 的 SageMaker AI 分散式執行期重新導向至訓練指令碼的 `main()` 函式。

```
def _mp_fn(index):
    main()
```

此函式會接受 `index` 引數，以指出叢集中目前 GPU 的等級，藉此進行分散式訓練。若要尋找更多範例指令碼，請參閱 [Hugging Face 轉換器語言模型範例指令碼](https://github.com/huggingface/transformers/blob/v4.21.1/examples/pytorch/language-modeling)。

**適用於轉換器 v4.17 及之前的版本，並搭配 PyTorch v1.10.2 及之前的版本**

SageMaker Training Compiler 使用替代機制來啟動分散式訓練任務，您不必在訓練指令碼中進行任何修改。相對而言，SageMaker Training Compiler 會要求您將 SageMaker AI 分散式訓練啟動器指令碼傳遞至 `entry_point` 引數，並將訓練指令碼傳遞至 SageMaker AI Hugging Face 估算器中的 `hyperparameters` 引數。

#### 將 SageMaker Training Compiler 搭配使用 `Trainer` 的最佳實務
<a name="training-compiler-pytorch-models-transformers-trainer-best-practices"></a>
+ 設定 [transformers.TrainingArgument](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments) 時，務必將 `optim` 引數設定為 `adamw_torch_xla`，以確保您使用的是 SyncFree 最佳化工具。另請參閱 *Hugging Face 轉換器文件*中的[最佳化工具](https://huggingface.co/docs/transformers/v4.23.1/en/perf_train_gpu_one#optimizer)。
+ 確保資料處理管道的輸送量高於訓練輸送量。您可以調整 [transformers.TrainingArgument](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments) 類別中的 `dataloader_num_workers` 和 `preprocessing_num_workers` 引數，來實現此一操作。通常，這些數量需要大於或等於 GPU 數量，但小於 CPU 數量。

完成訓練指令碼的調整後，請繼續前往[使用 SageMaker Training Compiler 執行 PyTorch 訓練任務](training-compiler-enable-pytorch.md)。

### 直接使用 PyTorch 的大型語言模型 (無需 Hugging Face 轉換器訓練器 API)
<a name="training-compiler-pytorch-models-non-trainer"></a>

如有直接使用 PyTorch 的訓練指令碼，則需對 PyTorch 訓練指令碼進行其他變更，以實作 PyTorch/XLA。按照指示修改指令碼，以妥善設定 PyTorch/XLA 本原值。

**Topics**
+ [適用於單一 GPU 訓練](#training-compiler-pytorch-models-non-trainer-single-gpu)
+ [適用於分散式訓練](#training-compiler-pytorch-models-non-trainer-distributed)
+ [將 SageMaker Training Compiler 與 PyTorch/XLA 搭配使用的最佳實務](#training-compiler-pytorch-models-best-practices)

#### 適用於單一 GPU 訓練
<a name="training-compiler-pytorch-models-non-trainer-single-gpu"></a>

1. 匯入最佳化程式庫。

   ```
   import torch_xla
   import torch_xla.core.xla_model as xm
   ```

1. 將目標裝置變更為 XLA 而非 `torch.device("cuda")`

   ```
   device=xm.xla_device()
   ```

1. 如果您使用 PyTorch 的 [自動混合精確度](https://pytorch.org/docs/stable/amp.html) (AMP)，請執行下列動作：

   1. 將 `torch.cuda.amp` 換成下列項目：

      ```
      import torch_xla.amp
      ```

   1. 將 `torch.optim.SGD` 和 `torch.optim.Adam` 取代為下列項目：

      ```
      import torch_xla.amp.syncfree.Adam as adam
      import torch_xla.amp.syncfree.SGD as SGD
      ```

   1. 將 `torch.cuda.amp.GradScaler` 換成下列項目：

      ```
      import torch_xla.amp.GradScaler as grad_scaler
      ```

1. 如果沒有使用 AMP，則請將 `optimizer.step()` 取代為下列項目：

   ```
   xm.optimizer_step(optimizer)
   ```

1. 如果您使用的是分散式資料載入器，請將資料載入器包裝在 PyTorch/XLA 的 `ParallelLoader` 類別中：

   ```
   import torch_xla.distributed.parallel_loader as pl
   parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
   ```

1. 若不使用 `parallel_loader` 時，請在訓練迴路的結尾新增 `mark_step`：

   ```
   xm.mark_step()
   ```

1. 若要檢查您的訓練，請使用 PyTorch/XLA 的模型檢查點方法：

   ```
   xm.save(model.state_dict(), path_to_save)
   ```

完成訓練指令碼的調整後，請繼續前往[使用 SageMaker Training Compiler 執行 PyTorch 訓練任務](training-compiler-enable-pytorch.md)。

#### 適用於分散式訓練
<a name="training-compiler-pytorch-models-non-trainer-distributed"></a>

除了上個 [適用於單一 GPU 訓練](#training-compiler-pytorch-models-non-trainer-single-gpu) 章節所列的變更內容之外，請新增下列變更項目，以便在 GPU 之間妥善分配工作負載。

1. 如果您使用了 AMP，請在 `scaler.scale(loss).backward()` 之後新增 `all_reduce`：

   ```
   gradients=xm._fetch_gradients(optimizer)
   xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
   ```

1. 如需為 `local_ranks` 和 `world_size` 設定變數，請使用與下列類似的程式碼：

   ```
   local_rank=xm.get_local_ordinal()
   world_size=xm.xrt_world_size()
   ```

1. 對於任何大於 `1` 的 `world_size` (`num_gpus_per_node*num_nodes`)，您必須定義一個訓練取樣器，該取樣器看起來應該類似於以下內容：

   ```
   import torch_xla.core.xla_model as xm
   
   if xm.xrt_world_size() > 1:
       train_sampler=torch.utils.data.distributed.DistributedSampler(
           train_dataset,
           num_replicas=xm.xrt_world_size(),
           rank=xm.get_ordinal(),
           shuffle=True
       )
   
   train_loader=torch.utils.data.DataLoader(
       train_dataset, 
       batch_size=args.batch_size,
       sampler=train_sampler,
       drop_last=args.drop_last,
       shuffle=False if train_sampler else True,
       num_workers=args.num_workers
   )
   ```

1. 進行下列變更，以確保您使用 `torch_xla distributed` 模組提供的 `parallel_loader`。

   ```
   import torch_xla.distributed.parallel_loader as pl
   train_device_loader=pl.MpDeviceLoader(train_loader, device)
   ```

   類似一般 PyTorch 載入器的 `train_device_loader` 函式如下所示：

   ```
   for step, (data, target) in enumerate(train_device_loader):
       optimizer.zero_grad()
       output=model(data)
       loss=torch.nn.NLLLoss(output, target)
       loss.backward()
   ```

   透過所有這些變更，您應該能夠在不使用轉換器訓練器 API 的情況下，使用任何 PyTorch 模型啟動分散式訓練。請注意，這些指示可用於單一節點多重 GPU 和多節點多重 GPU。

1. **適用於 PyTorch v1.11.0 及較新版本**

   若要使用 SageMaker Training Compiler 執行分散式訓練，您必須在訓練指令碼中新增下列 `_mp_fn()` 函式並包裝 `main()` 函式。它會將 `_mp_fn(index)` 函式呼叫從 PyTorch (`pytorchxla`) 的 SageMaker AI 分散式執行期重新導向至訓練指令碼的 `main()` 函式。

   ```
   def _mp_fn(index):
       main()
   ```

   此函式會接受 `index` 引數，以指出叢集中目前 GPU 的等級，藉此進行分散式訓練。若要尋找更多範例指令碼，請參閱 [Hugging Face 轉換器語言模型範例指令碼](https://github.com/huggingface/transformers/blob/v4.21.1/examples/pytorch/language-modeling)。

   **適用於轉換器 v4.17 及之前的版本，並搭配 PyTorch v1.10.2 及之前的版本**

   SageMaker Training Compiler 使用替代機制來啟動分散式訓練任務，並要求您將 SageMaker AI 分散式訓練啟動器指令碼傳遞至 `entry_point` 引數，並將訓練指令碼傳遞至 SageMaker AI Hugging Face 估算器中的 `hyperparameters` 引數。

完成訓練指令碼的調整後，請繼續前往[使用 SageMaker Training Compiler 執行 PyTorch 訓練任務](training-compiler-enable-pytorch.md)。

#### 將 SageMaker Training Compiler 與 PyTorch/XLA 搭配使用的最佳實務
<a name="training-compiler-pytorch-models-best-practices"></a>

如果您想要在原生 PyTorch 訓練指令碼中運用 SageMaker Training Compiler，您可能須先熟悉 [XLA 裝置上的 PyTorch](https://pytorch.org/xla/release/1.9/index.html)。以下各節列出了一些為 PyTorch 啟用 XLA 的最佳實務。

**注意**  
本節的最佳實務假設您會使用下列 PyTorch/XLA 模組：  

```
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
```

##### 了解 PyTorch/XLA 中的寬鬆模式
<a name="training-compiler-pytorch-models-best-practices-lazy-mode"></a>

PyTorch/XLA 和原生 PyTorch 之間的一大顯著差異為，PyTorch/XLA 系統在寬鬆模式中運作，而原生 PyTorch 則以嚴格模式執行。寬鬆模式中的張量會用來建置運算圖的預留位置，直到編譯和評估完成後其具體化為止。當您呼叫 PyTorch API 以使用張量和運算符建置運算內容時，PyTorch/XLA 系統將即時建置運算圖。當 `pl.MpDeviceLoader/pl.ParallelLoader` 明確或隱含地呼叫 `xm.mark_step()`，或者當您明確請求張量值 (例如透過呼叫 `loss.item()` 或 `print(loss)`) 時，運算圖會進行編譯並執行。

##### 使用 `pl.MpDeviceLoader/pl.ParallelLoader` 和 `xm.step_closure` 最小化 *compilation-and-executions* 的數量
<a name="training-compiler-pytorch-models-best-practices-minimize-comp-exec"></a>

為了獲得最佳效能，請記得啟動 *compilation-and-executions* 的可能方法，如[了解 PyTorch/XLA 中的寬鬆模式](#training-compiler-pytorch-models-best-practices-lazy-mode)中所述，並應盡量最小化 compilation-and-executions 的數量。理想情況下，每個訓練重複程序只需一個 compilation-and-execution，並由 `pl.MpDeviceLoader/pl.ParallelLoader` 自動初始化。`MpDeviceLoader` 已針對 XLA 進行最佳化，若可能的話，應時時使用以獲得最佳效能。在訓練期間，您可能需要檢查一些中繼結果，例如損失值。在這種情況下，應使用 `xm.add_step_closure()` 包裝寬鬆張量的列印，以避免不必要的 compilation-and-executions。

##### 使用 AMP 和 `syncfree` 最佳化工具
<a name="training-compiler-pytorch-models-best-practices-amp-optimizers"></a>

在自動混合精確度 (AMP) 模式中進行訓練，並運用 NVIDIA GPU 的 Tensor 核心，大幅加快您的訓練速度。SageMaker Training Compiler 提供針對 XLA 最佳化的 `syncfree` 最佳化工具，以提升 AMP 效能。目前，以下三個 `syncfree` 最佳化工具可以使用，如果可能的話，應使用這些工具以獲得最佳效能。

```
torch_xla.amp.syncfree.SGD
torch_xla.amp.syncfree.Adam
torch_xla.amp.syncfree.AdamW
```

這些 `syncfree` 最佳化工具應與漸層擴展/取消擴展搭配 `torch_xla.amp.GradScaler` 使用。

**提示**  
從 PyTorch 1.13.1 開始，SageMaker Training Compiler 可讓 PyTorch/XLA 在 `torch.optim` 或 `transformers.optimization` 當中自動覆寫最佳化工具 (例如 SGD、Adam、AdamW)，同時在 `torch_xla.amp.syncfree` (例如 `torch_xla.amp.syncfree.SGD`、`torch_xla.amp.syncfree.Adam`、`torch_xla.amp.syncfree.AdamW`) 使用它們本身的 syncfree 版本。您不需要變更您在訓練指令碼中定義最佳化工具的程式碼行。

# TensorFlow
<a name="training-compiler-tensorflow"></a>

攜帶您自己的 TensorFlow 模型至 SageMaker AI，並使用 SageMaker Training Compiler 執行訓練任務。

## TensorFlow 模型
<a name="training-compiler-tensorflow-models"></a>

SageMaker Training Compiler 會自動最佳化建置在原生 TensorFlow API 或高階 Keras API 上的模型訓練工作負載。

**提示**  
若要預先處理輸入資料集，請確保您使用靜態輸入形狀。動態輸入形狀可以啟動模型的重新編譯，並可能增加總訓練時間。

### 使用 Keras (建議)
<a name="training-compiler-tensorflow-models-keras"></a>

為了最佳的編譯器加速，我們建議您使用 TensorFlow Keras ([tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model)) 子類別的模型。

#### 適用於單一 GPU 訓練
<a name="training-compiler-tensorflow-models-keras-single-gpu"></a>

您不需要在訓練指令碼中進行其他變更。

### 不使用 Keras
<a name="training-compiler-tensorflow-models-no-keras"></a>

SageMaker Training Compiler 不支援 TensorFlow 的快速執行。因此，您應該使用 TensorFlow 函式裝飾項目 (`@tf.function`) 來包裝模型和訓練迴路，以利用編譯器加速。

SageMaker Training Compiler 執行圖形層級最佳化，並使用裝飾器來確保您的 TensorFlow 函式設定為以[圖形模式](https://www.tensorflow.org/guide/intro_to_graphs)執行。

#### 適用於單一 GPU 訓練
<a name="training-compiler-tensorflow-models-no-keras-single-gpu"></a>

TensorFlow 2.0 或更新版本預設為啟用快速執行，因此您應該在用於建構 TensorFlow 模型的每個函式前面加入 `@tf.function` 裝飾器。

## 配備 Hugging Face 轉換器的 TensorFlow 模型
<a name="training-compiler-tensorflow-models-transformers"></a>

配備 [Hugging Face 轉換器](https://huggingface.co/docs/transformers/index)的 TensorFlow 模型基於 TensorFlow 的 [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) API。Hugging Face 轉換器也為 TensorFlow 提供預先訓練的模型類別，有助於減少設定自然語言處理 (NLP) 模型的工作量。使用轉換器程式庫建立自己的訓練指令碼後，您可以使用 SageMaker AI `HuggingFace` 估算器搭配 SageMaker Training Compiler 組態類別執行訓練指令碼，如 [使用 SageMaker Training Compiler 執行 TensorFlow 訓練任務](training-compiler-enable-tensorflow.md) 的上一個主題所示。

SageMaker Training Compiler 會自動最佳化建置在原生 TensorFlow API 或高階 Keras API (例如 TensorFlow 轉換器模型) 的模型訓練工作負載。

**提示**  
當您在訓練指令碼中使用轉換器為 NLP 模型建立權杖化工具時，請確保您透過指定 `padding='max_length'` 來使用靜態輸入張量形狀。請勿使用 `padding='longest'`，因為填補至批次中最長的序列可能會變更每個訓練批次的張量形狀。動態輸入形狀可啟動模型的重新編譯，並可能增加總訓練時間。如需轉換器權杖化工具選項的更多相關資訊，請參閱 *Hugging Face 轉換器 文件*中的[填補和截斷](https://huggingface.co/docs/transformers/pad_truncation)。

**Topics**
+ [使用 Keras](#training-compiler-tensorflow-models-transformers-keras)
+ [不使用 Keras](#training-compiler-tensorflow-models-transformers-no-keras)

### 使用 Keras
<a name="training-compiler-tensorflow-models-transformers-keras"></a>

為了最佳的編譯器加速，我們建議您使用 TensorFlow Keras ([tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model)) 子類別的模型。正如 *Hugging Face 轉換器文件*中的[快速導覽](https://huggingface.co/docs/transformers/quicktour)頁面所述，您可以使用模型做為一般 TensorFlow Keras 模型。

#### 適用於單一 GPU 訓練
<a name="training-compiler-tensorflow-models-transformers-keras-single-gpu"></a>

您不需要在訓練指令碼中進行其他變更。

#### 適用於分散式訓練
<a name="training-compiler-tensorflow-models-transformers-keras-distributed"></a>

在 [https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) 呼叫範圍內使用 Keras API 建構和訓練模型時，SageMaker Training Compiler 加速可透明針對多重 GPU 工作負載運作。

1. 選擇正確的分散式訓練策略。

   1. 對於單一節點多重 GPU，請使用 `tf.distribute.MirroredStrategy` 設定策略。

      ```
      strategy = tf.distribute.MirroredStrategy()
      ```

   1. 對於多重節點多重 GPU，請在建立策略之前新增下列程式碼，以正確設定 TensorFlow 分散式訓練組態。

      ```
      def set_sm_dist_config():
          DEFAULT_PORT = '8890'
          DEFAULT_CONFIG_FILE = '/opt/ml/input/config/resourceconfig.json'
          with open(DEFAULT_CONFIG_FILE) as f:
              config = json.loads(f.read())
              current_host = config['current_host']
          tf_config = {
              'cluster': {
                  'worker': []
              },
              'task': {'type': 'worker', 'index': -1}
          }
          for i, host in enumerate(config['hosts']):
              tf_config['cluster']['worker'].append("%s:%s" % (host, DEFAULT_PORT))
              if current_host == host:
                  tf_config['task']['index'] = i
          os.environ['TF_CONFIG'] = json.dumps(tf_config)
      
      set_sm_dist_config()
      ```

       使用 `tf.distribute.MultiWorkerMirroredStrategy` 設定策略。

      ```
      strategy = tf.distribute.MultiWorkerMirroredStrategy()
      ```

1. 使用您選擇的策略，包裝模型。

   ```
   with strategy.scope():
       # create a model and do fit
   ```

### 不使用 Keras
<a name="training-compiler-tensorflow-models-transformers-no-keras"></a>

如果您想要使用沒有 Keras 的 TensorFlow 來具有自訂訓練迴路的自訂模型，您應該使用 TensorFlow 函式裝飾器 (`@tf.function`) 來包裝模型和訓練迴路，以利用編譯器加速。

SageMaker Training Compiler 執行圖形層級最佳化，並使用裝飾器來確保您的 TensorFlow 函式設定為以圖形模式執行。

#### 適用於單一 GPU 訓練
<a name="training-compiler-tensorflow-models-transformers-no-keras-single-gpu"></a>

TensorFlow 2.0 或更新版本預設為啟用快速執行，因此您應該在用於建構 TensorFlow 模型的每個函式前面加入 `@tf.function` 裝飾器。

#### 適用於分散式訓練
<a name="training-compiler-tensorflow-models-transformers-no-keras-distributed"></a>

除了[使用 Keras 進行分散式訓練](https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-tensorflow-models.html#training-compiler-tensorflow-models-transformers-keras)所需的變更之外，您還需要確保在每個 GPU 上執行的功能均已註釋 `@tf.function`，但不會註釋跨 GPU 通訊功能。範例訓練程式碼看起來應該如下所示：

```
@tf.function()
def compiled_step(inputs, outputs):
    with tf.GradientTape() as tape:
        pred=model(inputs, training=True)
        total_loss=loss_object(outputs, pred)/args.batch_size
    gradients=tape.gradient(total_loss, model.trainable_variables)
    return total_loss, pred, gradients

def train_step(inputs, outputs):
    total_loss, pred, gradients=compiled_step(inputs, outputs)
    if args.weight_decay > 0.:
        gradients=[g+v*args.weight_decay for g,v in zip(gradients, model.trainable_variables)]

    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss.update_state(total_loss)
    train_accuracy.update_state(outputs, pred)

@tf.function()
def train_step_dist(inputs, outputs):
    strategy.run(train_step, args= (inputs, outputs))
```

請注意，此指示可用於單一節點多重 GPU 和多重節點多重 GPU。