1. 總覽
在這個研究室中，您會使用 Vertex AI，透過 TensorFlow 在 Vertex AI 訓練中執行分散式訓練工作。
這個研究室是「Prototype to Production」系列影片的一部分。試用本研究室前，請務必完成先前的研究室。歡迎觀看相關系列影片，瞭解更多資訊：
。
課程內容
您將學習下列內容：
- 在搭載多個 GPU 的單一機器上執行分散式訓練
- 在多個機器中執行分散式訓練
在 Google Cloud 中執行這個研究室的總費用約為 $2 美元。
2. Vertex AI 簡介
這個研究室使用 Google Cloud 最新的 AI 產品服務。Vertex AI 將 Google Cloud 中的機器學習產品整合到流暢的開發體驗中。先前使用 AutoML 訓練的模型和自訂模型，都能透過不同的服務存取。這項新產品會與其他新產品一起合併為一個 API。您也可以將現有專案遷移至 Vertex AI。
Vertex AI 提供許多不同的產品，可支援端對端機器學習工作流程。本研究室將著重介紹下列產品：訓練和 Workbench
3. 分散式訓練總覽
如果您使用單一 GPU，TensorFlow 會使用這個加速器來加快模型訓練速度，您不須執行額外操作。不過，如果希望使用多個 GPU 提升額外效能，就需要使用 tf.distribute，這是在多部裝置上執行運算的 TensorFlow 模組。
本研究室的第一個部分使用的是 tf.distribute.MirroredStrategy，您只要稍微修改程式碼，就能將這個程式碼新增至訓練應用程式。這項策略會為電腦上的每個 GPU 建立模型副本。後續的漸層更新則會以同步方式進行。也就是說，每個 GPU 都會根據輸入資料的不同部分，計算向前和向後傳遞模型。接著，根據每個配量計算出的梯度，則會匯總所有 GPU 的梯度，並在名為 all-reduce 的處理程序中算出平均值。模型參數使用這些平均值的梯度更新。
研究室結束時的選用部分會使用 tf.distribute.MultiWorkerMirroredStrategy，這類似於 MirroredStrategy，差別在於其適用多部機器。這些機器也可能有多個 GPU。就像 MirroredStrategy、MultiWorkerMirroredStrategy 一樣，都是同步資料平行處理策略，只需變更少量程式碼即可使用。在一部機器上從同步資料平行處理移至多部機器的主要差異，在於現在每個步驟的結尾梯度都必須同步處理至同一機器中的所有 GPU 和叢集中所有機器。
您不需要具備詳細資料即可完成本研究室活動，但如果你想進一步瞭解分散式訓練在 TensorFlow 中的運作方式，請觀看以下影片：
4. 設定環境
完成「使用 Vertex AI 訓練自訂模型」研究室中的步驟，完成環境設定作業。
5. 單一機器、多 GPU 訓練
您必須將訓練應用程式程式碼放入 Docker 容器，並將這個容器推送至 Google Artifact Registry，藉此將分散式訓練工作提交至 Vertex AI。透過這種做法，您可以訓練以任何架構建構的模型。
首先，請在先前研究室建立的 Workbench 筆記本的啟動器選單中，開啟終端機視窗。
步驟 1：編寫訓練程式碼
建立名為 flowers-multi-gpu 的新目錄，並使用 cd 加入該目錄：
mkdir flowers-multi-gpu
cd flowers-multi-gpu
執行下列指令，為訓練程式碼建立目錄，以及建立以下程式碼的 Python 檔案。
mkdir trainer
touch trainer/task.py
現在 flowers-multi-gpu/ 目錄中應會顯示以下內容：
+ trainer/
+ task.py
接著，開啟剛剛建立的 task.py 檔案，並複製以下程式碼。
請將 BUCKET_ROOT 中的 {your-gcs-bucket} 替換成您在 Lab 1 中儲存花卉資料集的 Cloud Storage 值區。
import tensorflow as tf
import numpy as np
import os
## Replace {your-gcs-bucket} !!
BUCKET_ROOT='/gcs/{your-gcs-bucket}'
# Define variables
NUM_CLASSES = 5
EPOCHS=10
BATCH_SIZE = 32
IMG_HEIGHT = 180
IMG_WIDTH = 180
DATA_DIR = f'{BUCKET_ROOT}/flower_photos'
def create_datasets(data_dir, batch_size):
'''Creates train and validation datasets.'''
train_dataset = tf.keras.utils.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="training",
seed=123,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=batch_size)
validation_dataset = tf.keras.utils.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="validation",
seed=123,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=batch_size)
train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
validation_dataset = validation_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
return train_dataset, validation_dataset
def create_model():
'''Creates model.'''
model = tf.keras.Sequential([
tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
])
return model
def main():
# Create distribution strategy
strategy = tf.distribute.MirroredStrategy()
# Get data
GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)
# Wrap model creation and compilation within scope of strategy
with strategy.scope():
model = create_model()
model.compile(optimizer=tf.keras.optimizers.Adam(),
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=['accuracy'])
history = model.fit(
train_dataset,
validation_data=validation_dataset,
epochs=EPOCHS
)
model.save(f'{BUCKET_ROOT}/model_output')
if __name__ == "__main__":
main()
建構容器之前，讓我們先進一步瞭解程式碼。有一些元件專門使用分散式訓練。
- 在
main()函式中建立MirroredStrategy物件。接下來，您要將模型變數建立作業納入策略範圍內。這個步驟會告知 TensorFlow 應在 GPU 上鏡像哪些變數。 - 批量會由
num_replicas_in_sync向上擴充。在 TensorFlow 中使用同步資料平行處理策略時，調整批次大小是最佳做法。如要瞭解詳情，請按這裡。
步驟 2：建立 Dockerfile
如要將程式碼容器化，您必須建立 Dockerfile。而 Dockerfile 包含執行映像檔所需的所有指令。其會安裝所有必要的程式庫，並設定訓練程式碼的進入點。
在終端機中，於花卉目錄的根目錄中建立空白的 Dockerfile：
touch Dockerfile
現在 flowers-multi-gpu/ 目錄中應會顯示以下內容：
+ Dockerfile
+ trainer/
+ task.py
開啟 Dockerfile，並將下列內容複製到該檔案中：
FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-8
WORKDIR /
# Copies the trainer code to the docker image.
COPY trainer /trainer
# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]
步驟 3：建構容器
在終端機執行下列指令，定義專案的環境變數；請務必將 your-cloud-project 替換為您的專案 ID：
PROJECT_ID='your-cloud-project'
在 Artifact Registry 中建立存放區。我們會使用在第一個研究室中建立的存放區
REPO_NAME='flower-app'
使用 Artifact Registry 中容器映像檔的 URI 定義變數：
IMAGE_URI=us-central1-docker.pkg.dev/$PROJECT_ID/$REPO_NAME/flower_image_distributed:single_machine
設定 Docker
gcloud auth configure-docker \
us-central1-docker.pkg.dev
接著，從 flowers-multi-gpu 目錄的根目錄執行下列指令，以建構容器：
docker build ./ -t $IMAGE_URI
最後，將其推送至 Artifact Registry：
docker push $IMAGE_URI
將容器推送至 Artifact Registry 後，您就能開始執行訓練工作。
步驟 4：使用 SDK 執行工作
在本節中，您將瞭解如何使用 Vertex AI Python SDK 設定及啟動分散式訓練工作。
在 Launcher 中建立 TensorFlow 2 筆記本。
匯入 Vertex AI SDK。
from google.cloud import aiplatform
然後定義 CustomContainerTrainingJob。
您需要替換 container_uri 和 staging_bucket 中的 {YOUR_BUCKET} 中的 {PROJECT_ID}。
job = aiplatform.CustomContainerTrainingJob(display_name='flowers-multi-gpu',
container_uri='us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image_distributed:single_machine',
staging_bucket='gs://{YOUR_BUCKET}')
定義工作後，您就可以執行工作。將加速器數量設為 2。如果只使用 1 個 GPU，這「不會」視為分散式訓練。使用 2 個以上的加速器時，在同一部機器上安排分散式訓練。
my_custom_job.run(replica_count=1,
machine_type='n1-standard-4',
accelerator_type='NVIDIA_TESLA_V100',
accelerator_count=2)
您可以在控制台中查看工作的進度。
6. [選用] 多人培訓
透過採用多個 GPU 的單一機器進行分散式訓練後，現在您可以在多部機器上進行訓練，讓分散式訓練技能更上層樓。為降低費用，我們不會為這些機器新增任何 GPU，但您可以視需求新增 GPU 進行實驗。
在您的筆記本執行個體中開啟新的終端機視窗：
步驟 1：編寫訓練程式碼
建立名為 flowers-multi-machine 的新目錄，並使用 cd 加入該目錄：
mkdir flowers-multi-machine
cd flowers-multi-machine
執行下列指令，為訓練程式碼建立目錄，以及建立以下程式碼的 Python 檔案。
mkdir trainer
touch trainer/task.py
現在 flowers-multi-machine/ 目錄中應會顯示以下內容：
+ trainer/
+ task.py
接著，開啟剛剛建立的 task.py 檔案，並複製以下程式碼。
請將 BUCKET_ROOT 中的 {your-gcs-bucket} 替換成您在 Lab 1 中儲存花卉資料集的 Cloud Storage 值區。
import tensorflow as tf
import numpy as np
import os
## Replace {your-gcs-bucket} !!
BUCKET_ROOT='/gcs/{your-gcs-bucket}'
# Define variables
NUM_CLASSES = 5
EPOCHS=10
BATCH_SIZE = 32
IMG_HEIGHT = 180
IMG_WIDTH = 180
DATA_DIR = f'{BUCKET_ROOT}/flower_photos'
SAVE_MODEL_DIR = f'{BUCKET_ROOT}/multi-machine-output'
def create_datasets(data_dir, batch_size):
'''Creates train and validation datasets.'''
train_dataset = tf.keras.utils.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="training",
seed=123,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=batch_size)
validation_dataset = tf.keras.utils.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="validation",
seed=123,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=batch_size)
train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
validation_dataset = validation_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
return train_dataset, validation_dataset
def create_model():
'''Creates model.'''
model = tf.keras.Sequential([
tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
])
return model
def _is_chief(task_type, task_id):
'''Helper function. Determines if machine is chief.'''
return task_type == 'chief'
def _get_temp_dir(dirpath, task_id):
'''Helper function. Gets temporary directory for saving model.'''
base_dirpath = 'workertemp_' + str(task_id)
temp_dir = os.path.join(dirpath, base_dirpath)
tf.io.gfile.makedirs(temp_dir)
return temp_dir
def write_filepath(filepath, task_type, task_id):
'''Helper function. Gets filepath to save model.'''
dirpath = os.path.dirname(filepath)
base = os.path.basename(filepath)
if not _is_chief(task_type, task_id):
dirpath = _get_temp_dir(dirpath, task_id)
return os.path.join(dirpath, base)
def main():
# Create distribution strategy
strategy = tf.distribute.MultiWorkerMirroredStrategy()
# Get data
GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)
# Wrap variable creation within strategy scope
with strategy.scope():
model = create_model()
model.compile(optimizer=tf.keras.optimizers.Adam(),
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=['accuracy'])
history = model.fit(
train_dataset,
validation_data=validation_dataset,
epochs=EPOCHS
)
# Determine type and task of the machine from
# the strategy cluster resolver
task_type, task_id = (strategy.cluster_resolver.task_type,
strategy.cluster_resolver.task_id)
# Based on the type and task, write to the desired model path
write_model_path = write_filepath(SAVE_MODEL_DIR, task_type, task_id)
model.save(write_model_path)
if __name__ == "__main__":
main()
建構容器之前，讓我們先進一步瞭解程式碼。程式碼必須包含幾個元件，訓練應用程式才能使用 MultiWorkerMirroredStrategy。
- 在
main()函式中建立MultiWorkerMirroredStrategy物件。接下來，您要將模型變數建立作業納入策略範圍內。這個關鍵步驟會告知 TensorFlow 應在備用資源間鏡像哪些變數。 - 批量會由
num_replicas_in_sync向上擴充。在 TensorFlow 中使用同步資料平行處理策略時，調整批次大小是最佳做法。 - 在需要多工作站的情況下，模型的儲存方式會稍微複雜一點，因為每個工作站的目的地都必須不同。主要工作站將儲存至所需模型目錄，其他 worker 則會將模型儲存至臨時目錄。這些臨時目錄不能重複，以免多個工作站寫入同一個位置。保存能夠包括集體作業，意味著所有工作人員必須儲存，而不只是主要業務。
_is_chief()、_get_temp_dir()、write_filepath()和main()函式全都包含有助於儲存模型的樣板程式碼。
步驟 2：建立 Dockerfile
如要將程式碼容器化，您必須建立 Dockerfile。而 Dockerfile 包含執行映像檔所需的所有指令。其會安裝所有必要的程式庫，並設定訓練程式碼的進入點。
在終端機中，於花卉目錄的根目錄中建立空白的 Dockerfile：
touch Dockerfile
現在 flowers-multi-machine/ 目錄中應會顯示以下內容：
+ Dockerfile
+ trainer/
+ task.py
開啟 Dockerfile，並將下列內容複製到該檔案中：
FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-8
WORKDIR /
# Copies the trainer code to the docker image.
COPY trainer /trainer
# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]
步驟 3：建構容器
在終端機執行下列指令，定義專案的環境變數；請務必將 your-cloud-project 替換為您的專案 ID：
PROJECT_ID='your-cloud-project'
在 Artifact Registry 中建立存放區。我們會使用在第一個研究室中建立的存放區
REPO_NAME='flower-app'
使用 Google Artifact Registry 中容器映像檔的 URI 定義變數：
IMAGE_URI=us-central1-docker.pkg.dev/$PROJECT_ID/$REPO_NAME/flower_image_distributed:multi_machine
設定 Docker
gcloud auth configure-docker \
us-central1-docker.pkg.dev
接著，從 flowers-multi-machine 目錄的根目錄執行下列指令，以建構容器：
docker build ./ -t $IMAGE_URI
最後，將其推送至 Artifact Registry：
docker push $IMAGE_URI
將容器推送至 Artifact Registry 後，您就能開始執行訓練工作。
步驟 4：使用 SDK 執行工作
在本節中，您將瞭解如何使用 Vertex AI Python SDK 設定及啟動分散式訓練工作。
在 Launcher 中建立 TensorFlow 2 筆記本。
匯入 Vertex AI SDK。
from google.cloud import aiplatform
接著定義 worker_pool_specs。
Vertex AI 提供 4 個工作站集區，可處理不同類型的機器工作。
工作站集區 0 會設定主要、主要、排程器或「主要執行個體」。在 MultiWorkerMirroredStrategy 中，所有機器都會指派為工作站，這些機器就是要執行複製運算的實體機器。除了每部機器都是工作站，也需要一個工作站，以便完成一些額外的工作，例如儲存查核點以及將摘要檔案寫入 TensorBoard。這部機器稱為首領。由於工作站集區只有一個主要工作站，工作站集區 0 的工作站數量永遠會是 1。
工作站集區 1 可讓您為叢集設定額外的工作站。
worker_pool_specs 清單中的第一個字典代表工作站集區 0，第二個字典則代表工作站集區 1。這個範例中的兩項設定完全相同。不過，如果您想跨 3 部機器進行訓練，則必須將 replica_count 設為 2，在工作站集區 1 中新增額外的工作站。如要新增 GPU，您必須在兩個工作站集區的 machine_spec 中新增 accelerator_type 和 accelerator_count 引數。請注意，如要使用 MultiWorkerMirroredStrategy 搭配 GPU，叢集中每部機器的 GPU 數量必須相同。否則工作就會失敗。
您需要替換 image_uri 中的 {PROJECT_ID}。
# The spec of the worker pools including machine type and Docker image
# Be sure to replace PROJECT_ID in the "image_uri" with your project.
worker_pool_specs=[
{
"replica_count": 1,
"machine_spec": {
"machine_type": "n1-standard-4",
},
"container_spec": {"image_uri": "us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image_distributed:multi_machine"}
},
{
"replica_count": 1,
"machine_spec": {
"machine_type": "n1-standard-4",
},
"container_spec": {"image_uri": "us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image_distributed:multi_machine"}
}
]
接著建立並執行 CustomJob，並將 staging_bucket 中的 {YOUR_BUCKET} 替換為專案中的值區，以便執行測試作業。
my_custom_job = aiplatform.CustomJob(display_name='flowers-multi-worker',
worker_pool_specs=worker_pool_specs,
staging_bucket='gs://{YOUR_BUCKET}')
my_custom_job.run()
您可以在控制台中查看工作的進度。
🎉 恭喜！🎉
您已瞭解如何使用 Vertex AI 執行下列作業：
- 使用 TensorFlow 執行分散式訓練工作
如要進一步瞭解 Vertex 的其他部分，請參閱說明文件。
7. 清除
因為我們設定了筆記本在 60 分鐘閒置後逾時，所以我們不必擔心如何關閉執行個體。如要手動關閉執行個體，請前往控制台的「Vertex AI Workbench」專區，然後按一下「Stop」按鈕。如想完全刪除筆記本，請按一下「刪除」按鈕。
如要刪除 Storage 值區，請使用 Cloud 控制台中的導覽選單前往「Storage」(儲存空間)、選取值區，然後點選「Delete」(刪除)：
