<jupyter_start><jupyter_text>Distributed Training with GPUs on Cloud AI Platform**Learning Objectives:**  1. Setting up the environment  1. Create a model to train locally  1. Train on multiple GPUs/CPUs with MultiWorkerMirrored StrategyIn this notebook, we will walk through using Cloud AI Platform to perform distributed training using the `MirroredStrategy` found within `tf.keras`. This strategy will allow us to use the synchronous AllReduce strategy on a VM with multiple GPUs attached.Each learning objective will correspond to a __TODO__ in this student lab notebook -- try to complete this notebook first and then review the [Solution Notebook](https://github.com/GoogleCloudPlatform/training-data-analyst/blob/master/courses/machine_learning/deepdive2/production_ml/solutions/distributed_training.ipynb) for reference.<jupyter_code>!sudo chown -R jupyter:jupyter /home/jupyter/training-data-analyst<jupyter_output><empty_output><jupyter_text>Next we will configure our environment. Be sure to change the `PROJECT_ID` variable in the below cell to your Project ID. This will be the project to which the Cloud AI Platform resources will be billed. We will also create a bucket for our training artifacts (if it does not already exist).  Lab Task 1: Setting up the environment<jupyter_code>import os
# TODO 1
PROJECT_ID = "qwiklabs-gcp-03-3fcd4f8a7fba"  # Replace with your PROJECT
BUCKET = PROJECT_ID 
REGION = 'australia-southeast1'
os.environ["PROJECT_ID"] = PROJECT_ID
os.environ["BUCKET"] = BUCKET<jupyter_output><empty_output><jupyter_text>Since we are going to submit our training job to Cloud AI Platform, we need to create our trainer package. We will create the `train` directory for our package and create a blank `__init__.py` file so Python knows that this folder contains a package.<jupyter_code>!mkdir train
!touch train/__init__.py<jupyter_output><empty_output><jupyter_text>Next we will create a module containing a function which will create our model. Note that we will be using the Fashion MNIST dataset. Since it's a small dataset, we will simply load it into memory for getting the parameters for our model.Our model will be a DNN with only dense layers, applying dropout to each hidden layer. We will also use ReLU activation for all hidden layers.<jupyter_code>%%writefile train/model_definition.py
import tensorflow as tf
import numpy as np

# Get data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

def create_model():
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Flatten(input_shape=x_train.shape[1:]))
    model.add(tf.keras.layers.Dense(1028))
    model.add(tf.keras.layers.Activation('relu'))
    model.add(tf.keras.layers.Dropout(0.5))
    model.add(tf.keras.layers.Dense(512))
    model.add(tf.keras.layers.Activation('relu'))
    model.add(tf.keras.layers.Dropout(0.5))
    model.add(tf.keras.layers.Dense(256))
    model.add(tf.keras.layers.Activation('relu'))
    model.add(tf.keras.layers.Dropout(0.5))
    model.add(tf.keras.layers.Dense(10))
    model.add(tf.keras.layers.Activation('softmax'))
    return model<jupyter_output>Writing train/model_definition.py<jupyter_text>Before we submit our training jobs to Cloud AI Platform, let's be sure our model runs locally. We will call the `model_definition` function to create our model and use `tf.keras.datasets.fashion_mnist.load_data()` to import the Fashion MNIST dataset.  Lab Task 2: Create a model to train locally<jupyter_code>import os
import time
import tensorflow as tf
import numpy as np
from train import model_definition

def create_dataset(X, Y, epochs, batch_size):
    dataset = tf.data.Dataset.from_tensor_slices((X, Y))
    return dataset.repeat(epochs).batch(batch_size, drop_remainder=True)
# Get data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

ds_train = create_dataset(x_train, y_train, 20, 5000)
ds_test = create_dataset(x_test, y_test, 1, 1000)
model = model_definition.create_model()

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    loss='sparse_categorical_crossentropy',
    metrics=['sparse_categorical_accuracy']
)
    
start = time.time()
model.fit(
    ds_train,
    validation_data=ds_test, 
    verbose=1
)
print(f"Training time without GPUs locally: {time.time() - start} sec.")<jupyter_output>240/240 [==============================] - 134s 558ms/step - loss: 5.0578 - sparse_categorical_accuracy: 0.6339 - val_loss: 0.6118 - val_sparse_categorical_accuracy: 0.8129
Training time without GPUs locally: 134.2877902984619<jupyter_text>Train on multiple GPUs/CPUs with MultiWorkerMirrored Strategy That took a few minutes to train our model for 20 epochs. Let's see how we can do better using Cloud AI Platform. We will be leveraging the `MultiWorkerMirroredStrategy` supplied in `tf.distribute`. The main difference between this code and the code from the local test is that we need to compile the model within the scope of the strategy. When we do this our training op will use information stored in the `TF_CONFIG` variable to assign ops to the various devices for the AllReduce strategy. After the training process finishes, we will print out the time spent training. Since it takes a few minutes to spin up the resources being used for training on Cloud AI Platform, and this time can vary, we want a consistent measure of how long training took.Note: When we train models on Cloud AI Platform, the `TF_CONFIG` variable is automatically set. So we do not need to worry about adjusting based on what cluster configuration we use.<jupyter_code>%%writefile train/train_mult_worker_mirrored.py

strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

# Get data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

ds_train = create_dataset(x_train, y_train, 20, 5000)
ds_test = create_dataset(x_test, y_test, 1, 1000)
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
with strategy.scope():
    model = model_definition.create_model()
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss='sparse_categorical_crossentropy',
        metrics=['sparse_categorical_accuracy']
    )
    
start = time.time()
model.fit(
    ds_train,
    validation_data=ds_test, 
    verbose=2
)
print(f"Training time with multiple GPUs: {time.time() - start} sec.")<jupyter_output>Writing train/train_mult_worker_mirrored.py<jupyter_text>Lab Task 3: Training with multiple GPUs/CPUs on created model using MultiWorkerMirrored Strategy First we will train a model without using GPUs to give us a baseline. We will use a consistent format throughout the trials. We will define a `config.yaml` file to contain our cluster configuration and the pass this file in as the value of a command-line argument `--config`.In our first example, we will use a single `n1-highcpu-16` VM.<jupyter_code>%%writefile config.yaml
# TODO 3a
# Configure a master worker
trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
%%bash

now=$(date +"%Y%m%d_%H%M%S")
JOB_NAME="cpu_only_fashion_minst_$now"

gcloud ai-platform jobs submit training $JOB_NAME \
  --staging-bucket=gs://$BUCKET \
  --package-path=train \
  --module-name=train.train_mult_worker_mirrored \
  --runtime-version=2.3 \
  --python-version=3.7 \
  --region=australia-southeast1 \
  --config config.yaml

# us-west1, rather than australia-southeast1<jupyter_output>jobId: cpu_only_fashion_minst_20220307_091331
state: QUEUED<jupyter_text>If we go through the logs, we see that the training job will take around 5-7 minutes to complete. Let's now attach two Nvidia Tesla K80 GPUs and rerun the training job.<jupyter_code>%%writefile config.yaml
# TODO 3b
trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    acceleratorConfig:
      count: 2
      type: NVIDIA_TESLA_K80
%%bash

now=$(date +"%Y%m%d_%H%M%S")
JOB_NAME="multi_gpu_fashion_minst_2gpu_$now"

gcloud ai-platform jobs submit training $JOB_NAME \
  --staging-bucket=gs://$BUCKET \
  --package-path=train \
  --module-name=train.train_mult_worker_mirrored \
  --runtime-version=2.3 \
  --python-version=3.7 \
  --region=us-west1 \
  --config config.yaml

# us-west1, rather than australia-southeast1<jupyter_output>jobId: multi_gpu_fashion_minst_2gpu_20220307_091536
state: QUEUED<jupyter_text>That was a lot faster! The training job will take upto 5-10 minutes to complete. Let's keep going and add more GPUs!<jupyter_code>%%writefile config.yaml
# TODO 3c
trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    acceleratorConfig:
      count: 4
      type: NVIDIA_TESLA_K80
%%bash

now=$(date +"%Y%m%d_%H%M%S")
JOB_NAME="multi_gpu_fashion_minst_4gpu_$now"

gcloud ai-platform jobs submit training $JOB_NAME \
  --staging-bucket=gs://$BUCKET \
  --package-path=train \
  --module-name=train.train_mult_worker_mirrored \
  --runtime-version=2.3 \
  --python-version=3.7 \
  --region=us-west1 \
  --config config.yaml

# us-west1, rather than australia-southeast1<jupyter_output>jobId: multi_gpu_fashion_minst_4gpu_20220307_091722
state: QUEUED
{'fraction_non_alphanumeric': 0.09547950992817913, 'fraction_numerical': 0.022496831432192648, 'mean_word_length': 4.925531914893617, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '471166', 'license': 'Apache-2.0', 'n_tokens_mistral': 3251, 'n_tokens_neox': 3091, 'n_words': 925, 'path': 'C07/C07-w02-Lab5_distributed_training.ipynb', 'repo_name': 'pascal-p/ML_wtih_TF_GCP', 'size': '9468'}