Distributed data parallel MNIST training with TensorFlow 2 and SageMaker Distributed
This notebook’s CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.
Amazon SageMaker’s distributed library can be used to train deep learning models faster and cheaper. The data parallel feature in this library (smdistributed.dataparallel
) is a distributed data parallel training framework for PyTorch, TensorFlow, and MXNet. This notebook demonstrates how to use the SageMaker distributed data library to train
a PyTorch model using the MNIST dataset.
This notebook example shows how to use smdistributed.dataparallel
with TensorFlow2 in SageMaker using MNIST dataset.
For more information:
NOTE: This example requires SageMaker Python SDK v2.X.
Dataset
This example uses the MNIST dataset. MNIST is a widely used dataset for handwritten digit classification. It consists of 70,000 labeled 28x28 pixel grayscale images of handwritten digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits).
[ ]:
pip install sagemaker --upgrade
SageMaker role
The following code cell defines role
which is the IAM role ARN used to create and run SageMaker training and hosting jobs. This is the same IAM role used to create this SageMaker Notebook instance.
role
must have permission to create a SageMaker training job and launch an endpoint to host a model. For granular policies you can use to grant these permissions, see Amazon SageMaker Roles. If you do not require fine-tuned permissions for this demo, you can used the IAM managed policy AmazonSageMakerFullAccess to complete this demo.
[ ]:
import sagemaker
sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()
To verify that the role above has required permissions:
Go to the IAM console: https://console.aws.amazon.com/iam/home.
Select Roles.
Enter the role name in the search box to search for that role.
Select the role.
Use the Permissions tab to verify this role has required permissions attached.
Model training with SageMaker distributed data parallel
Training script
The train_tensorflow_smdataparallel_mnist.py
script provides the code you need for training a SageMaker model using smdistributed.dataparallel
’s DistributedGradientTape
. The training script is very similar to a TensorFlow2 training script you might run outside of SageMaker, and has been modified to run with this library. smdistributed.dataparallel
’s TensorFlow client provides an alternative to native DistributedGradientTape
.
For details about how to use smdistributed.dataparallel
’s DDP in your native TensorFlow 2 script, see the Modify a TensorFlow 2.x Training Script Using SMD Data Parallel.
[ ]:
!pygmentize code/train_tensorflow_smdataparallel_mnist.py
SageMaker TensorFlow Estimator function options
In the following code block, you can update the estimator function to use a different instance type, instance count, and distrubtion strategy. You’re also passing in the training script you reviewed in the previous cell.
Instance types
smdistributed.dataparallel
supports model training on SageMaker with the following instance types only. For best performance, it is recommended you use an instance type that supports Amazon Elastic Fabric Adapter (ml.p3dn.24xlarge and ml.p4d.24xlarge).
ml.p3.16xlarge
ml.p3dn.24xlarge [Recommended]
ml.p4d.24xlarge [Recommended]
Instance count
To get the best performance and the most out of smdistributed.dataparallel
, you should use at least 2 instances, but you can also use 1 for testing this example.
Distribution strategy
Note that to use DDP mode, you update the the distribution
strategy, and set it to use smdistributed dataparallel
.
[ ]:
from sagemaker.tensorflow import TensorFlow
estimator = TensorFlow(
base_job_name="tensorflow2-smdataparallel-mnist",
source_dir="code",
entry_point="train_tensorflow_smdataparallel_mnist.py",
role=role,
py_version="py37",
framework_version="2.4.1",
# For training with multinode distributed training, set this count. Example: 2
instance_count=2,
# For training with p3dn instance use - ml.p3dn.24xlarge, with p4dn instance use - ml.p4d.24xlarge
instance_type="ml.p3.16xlarge",
sagemaker_session=sagemaker_session,
# Training using SMDataParallel Distributed Training Framework
distribution={"smdistributed": {"dataparallel": {"enabled": True}}},
)
[ ]:
estimator.fit()
Now that you have a trained model, you can deploy an endpoint to host the model. After you deploy the endpoint, you can then test it with inference requests. The following cell will store the model_data variable to be used with the inference notebook.
[ ]:
model_data = estimator.model_data
print("Storing {} as model_data".format(model_data))
%store model_data
[ ]:
predictor = estimator.deploy(initial_instance_count=1, instance_type="ml.m4.xlarge")
[ ]:
print(predictor.endpoint_name)
[ ]:
import tensorflow as tf
import numpy as np
(mnist_images, mnist_labels), _ = tf.keras.datasets.mnist.load_data(path="/tmp/data")
[ ]:
for i in range(10):
data = mnist_images[i].reshape(1, 28, 28, 1)
predict_response = predictor.predict(data)
print("========================================")
label = mnist_labels[i]
predict_label = np.argmax(predict_response["predictions"])
print("label is {}".format(label))
print("prediction is {}".format(predict_label))
Cleanup
If you don’t intend to try out inference or to do anything else with the endpoint, you should delete the endpoint.
[ ]:
predictor.delete_endpoint()
Notebook CI Test Results
This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.