github-featured/aws--amazon-sagemaker-examples
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Fix JAX example to use BYOS (#3723)

Browse Source 
* Fix JAX example to use BYOS

* Spelling mistakes

Co-authored-by: atqy <95724753+atqy@users.noreply.github.com>
This commit is contained in:
Sean Morgan
2022-12-28 13:40:21 -08:00
committed by GitHub
parent 6a1863372e
commit df6c76360e
6 changed files with 133 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
advanced_functionality/jax_bring_your_own/sagemaker_jax.py
-60
View File
@@ -1,60 +0,0 @@
# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""
Custom Framework Estimator for JAX
"""
from sagemaker.estimator import Framework
from sagemaker.tensorflow.model import TensorFlowModel
from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
class JaxEstimator(Framework):
def __init__(
self,
entry_point,
source_dir=None,
hyperparameters=None,
image_uri=None,
**kwargs
):
super(JaxEstimator, self).__init__(
entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
)
def create_model(
self,
role=None,
vpc_config_override=VPC_CONFIG_DEFAULT,
entry_point=None,
source_dir=None,
dependencies=None,
**kwargs
):
"""Creates ``TensorFlowModel`` object to be used for creating SageMaker model entities"""
kwargs["name"] = self._get_or_create_name(kwargs.get("name"))
if "enable_network_isolation" not in kwargs:
kwargs["enable_network_isolation"] = self.enable_network_isolation()
return TensorFlowModel(
model_data=self.model_data,
role=role or self.role,
container_log_level=self.container_log_level,
framework_version="2.3.1",
sagemaker_session=self.sagemaker_session,
vpc_config=self.get_vpc_config(vpc_config_override),
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
**kwargs
)
advanced_functionality/jax_bring_your_own/train_deploy_jax.ipynb
+90 -109
View File
@@ -5,9 +5,29 @@
"metadata": {},
"source": [
"# Training and Deploying ML Models using JAX on SageMaker\n",
"Amazon SageMaker provides you the flexibility to train models using any framework that can work in a Docker container. In this example we'll show how to utilize the Bring-Your-Own-Container (BYOC) paradigm to train machine learning models using the increasingly popular [JAX library from Google](https://github.com/google/jax). We'll train a fashion mnist classification model using vanilla JAX, another using `jax.experimental.stax`, and a final model using the [higher level Trax library from Google](https://github.com/google/trax).\n",
"Amazon SageMaker provides you the flexibility to train models using our pre-built machine learning containers or your own bespoke container. We'll refer to these strategies as Bring-Your-Own-Script **(BYOS)** and Bring-Your-Own-Container **(BYOC)** in this tutorial. \n",
"\n",
"For both of these demos, we'll show how both JAX and Trax can serialize models using the TensorFlow standard [SavedModel format](https://www.tensorflow.org/guide/saved_model). This enables us to train these models in a custom container, but then deploy them using the managed and optimized SageMaker TensorFlow inference containers.\n"
"### Bring Your Own JAX Script\n",
"\n",
"In this notebook, we'll show how to extend our optimized TensorFlow containers to train machine learning models using the increasingly popular [JAX library](https://github.com/google/jax). We'll train a fashion MNIST classification model using vanilla JAX, another using `jax.experimental.stax`, and a final model using the [higher level Trax library](https://github.com/google/trax).\n",
"\n",
"For all three patterns, we'll show how the JAX models can be serialized as standard TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model). This enables us to seamlessly deploy the models using the managed and optimized SageMaker TensorFlow inference containers.\n",
"\n",
"\n",
"### Bring Your Own JAX Container\n",
"\n",
"We've included a dockerfile in this repo directory to show how you can build your own bespoke JAX container with support for GPUs on SageMaker. Unfortunately, the NVIDIA/CUDA Dockerhub containers have a [deletion policy](https://gitlab.com/nvidia/container-images/cuda/blob/master/doc/support-policy.md), so we're unable to assert that the container can be built through time. Nonetheless, you can trivially adapt a newer version of the container if your workload requires a custom container. For more information on running BYOC on SageMaker see the [documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html).\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade sagemaker"
]
},
{
@@ -16,102 +36,22 @@
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import json\n",
"\n",
"import boto3\n",
"import sagemaker\n",
"from sagemaker import get_execution_role\n",
"from sagemaker.tensorflow import TensorFlow\n",
"\n",
"from sagemaker_jax import JaxEstimator\n",
"\n",
"client = boto3.client(\"sts\")\n",
"account = client.get_caller_identity()[\"Account\"]\n",
"role = get_execution_role()\n",
"my_session = boto3.session.Session()\n",
"region = my_session.region_name\n",
"\n",
"container_name = \"sagemaker-jax\"\n",
"ecr_image = \"{}.dkr.ecr.{}.amazonaws.com/{}\".format(account, region, container_name)"
"role = get_execution_role()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Custom Framework Estimator\n",
"Since we'll be saving our JAX and Trax models as SavedModel format, we can create a subclass of the base [SageMaker Framework estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html?highlight=Framework#sagemaker.estimator.Framework). This will enable us to specify a custom `create_model` method which leverages the existing TensorFlowModel class to launch inference containers"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pygmentize sagemaker_jax.py"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Docker Container\n",
"Our custom training container is straight forward, though there are a few things worth mentioning that can be seen in the comments"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!cat docker/Dockerfile"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Building and Publishing the Image\n",
"The below shell script must be run if the docker image has not already been pushed to the Elastic Container Registry. \n",
"## Installing JAX in SageMaker TensorFlow Containers\n",
"\n",
"**NOTE: Since SageMaker studio is already running inside a Docker container, this script cannot be run inside SageMaker Studio. Please push your container using awscli or use this toolkit: https://github.com/aws-samples/sagemaker-studio-image-build-cli**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# %%sh\n",
"When using BYOS with managed SageMaker containers, you can trivially install extra dependencies by providing a `requirements.txt` within the `source_dir` that contains your training scripts. At runtime these dependencies will be installed prior to executing the training script, so we can utilize our optimized TensorFlow GPU container to utilize JAX with CUDA support.\n",
"\n",
"# container_name=sagemaker-jax\n",
"# account=$(aws sts get-caller-identity --query Account --output text)\n",
"\n",
"# # Get the region defined in the current configuration (default to us-west-2 if none defined)\n",
"# region=$(aws configure get region)\n",
"# region=${region:-us-west-2}\n",
"\n",
"# fullname=\"${account}.dkr.ecr.${region}.amazonaws.com/${container_name}\"\n",
"\n",
"# # If the repository doesn't exist in ECR, create it.\n",
"# aws ecr describe-repositories --repository-names \"${container_name}\" > /dev/null 2>&1\n",
"# if [ $? -ne 0 ]\n",
"# then\n",
"# aws ecr create-repository --repository-name \"${container_name}\" > /dev/null\n",
"# fi\n",
"\n",
"# # Get the login command from ECR and execute it directly\n",
"# $(aws ecr get-login --region ${region} --no-include-email)\n",
"\n",
"# # Build the docker image locally with the image name and then push it to ECR\n",
"# # with the full name.\n",
"# docker build -t ${container_name} docker/\n",
"# docker tag ${container_name} ${fullname}\n",
"\n",
"# docker push ${fullname}"
"To be specific, any container that has the [sagemaker-training-toolkit](https://github.com/aws/sagemaker-training-toolkit) supports installing additional dependencies from `requirements.txt`\n"
]
},
{
@@ -119,19 +59,49 @@
"metadata": {},
"source": [
"## Serializing models as SavedModel format\n",
"In the upcoming training jobs we'll be training a vanilla JAX model, a Stax model, and a Trax model on the [fashion mnist dataset](https://github.com/zalandoresearch/fashion-mnist).\n",
"In the upcoming training jobs we'll be training a vanilla JAX model, a Stax model, and a Trax model on the [fashion MNIST dataset](https://github.com/zalandoresearch/fashion-mnist).\n",
"The full details of the model can be seen in the `training_scripts/` directory, but it is worth calling out the methods for serialization.\n",
"\n",
"The JAX model utilizes the new experimental jax2tf converter: https://github.com/google/jax/tree/master/jax/experimental/jax2tf\n",
"The JAX/Stax models utilize the new jax2tf converter: https://github.com/google/jax/tree/master/jax/experimental/jax2tf\n",
"\n",
"The Trax model utilizes the new trax2keras functionality: https://github.com/google/trax/blob/master/trax/trax2keras.py"
"```python\n",
"def save_model_tf(prediction_function, params_to_save):\n",
" tf_fun = jax2tf.convert(prediction_function, enable_xla=False)\n",
" param_vars = tf.nest.map_structure(lambda param: tf.Variable(param), params_to_save)\n",
"\n",
" tf_graph = tf.function(\n",
" lambda inputs: tf_fun(param_vars, inputs),\n",
" autograph=False,\n",
" jit_compile=False,\n",
" )\n",
"\n",
"```\n",
"\n",
"\n",
"The Trax model utilizes the new trax2keras functionality: https://github.com/google/trax/blob/master/trax/trax2keras.py\n",
"\n",
"```python\n",
"def save_model_tf(model_to_save):\n",
" \"\"\"\n",
" Serialize a TensorFlow graph from trained Trax Model\n",
" :param model_to_save: Trax Model\n",
" \"\"\"\n",
" keras_layer = trax.AsKeras(model_to_save, batch_size=1)\n",
" inputs = tf.keras.Input(shape=(28, 28, 1))\n",
" hidden = keras_layer(inputs)\n",
"\n",
" keras_model = tf.keras.Model(inputs=inputs, outputs=hidden)\n",
" keras_model.save(\"/opt/ml/model/1\", save_format=\"tf\")\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train using Vanilla JAX"
"## Train using Vanilla JAX\n",
"\n",
"Note: Our `source_dir` directory contains a `requirements.txt` that will install JAX with CUDA support"
]
},
{
@@ -140,14 +110,15 @@
"metadata": {},
"outputs": [],
"source": [
"vanilla_jax_estimator = JaxEstimator(\n",
" image_uri=ecr_image,\n",
"vanilla_jax_estimator = TensorFlow(\n",
" role=role,\n",
" instance_count=1,\n",
" base_job_name=container_name + \"-jax\",\n",
" base_job_name=\"jax\",\n",
" framework_version=\"2.10\",\n",
" py_version=\"py39\",\n",
" source_dir=\"training_scripts\",\n",
" entry_point=\"train_jax.py\",\n",
" instance_type=\"ml.p2.xlarge\",\n",
" instance_type=\"ml.p3.2xlarge\",\n",
" hyperparameters={\"num_epochs\": 3},\n",
")\n",
"vanilla_jax_estimator.fit(logs=False)"
@@ -166,14 +137,15 @@
"metadata": {},
"outputs": [],
"source": [
"stax_estimator = JaxEstimator(\n",
" image_uri=ecr_image,\n",
"stax_estimator = TensorFlow(\n",
" role=role,\n",
" instance_count=1,\n",
" base_job_name=container_name + \"-jax\",\n",
" base_job_name=\"stax\",\n",
" framework_version=\"2.10\",\n",
" py_version=\"py39\",\n",
" source_dir=\"training_scripts\",\n",
" entry_point=\"train_stax.py\",\n",
" instance_type=\"ml.p2.xlarge\",\n",
" instance_type=\"ml.p3.2xlarge\",\n",
" hyperparameters={\"num_epochs\": 3},\n",
")\n",
"\n",
@@ -193,17 +165,19 @@
"metadata": {},
"outputs": [],
"source": [
"trax_estimator = JaxEstimator(\n",
" image_uri=ecr_image,\n",
"trax_estimator = TensorFlow(\n",
" role=role,\n",
" instance_count=1,\n",
" base_job_name=container_name + \"-trax\",\n",
" base_job_name=\"trax\",\n",
" framework_version=\"2.10\",\n",
" py_version=\"py39\",\n",
" source_dir=\"training_scripts\",\n",
" entry_point=\"train_trax.py\",\n",
" instance_type=\"ml.p2.xlarge\",\n",
" instance_type=\"ml.p3.2xlarge\",\n",
" hyperparameters={\"train_steps\": 1000},\n",
")\n",
"\n",
"\n",
"trax_estimator.fit(logs=False)"
]
},
@@ -211,8 +185,8 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Deploy Both Models to prebuilt TF Containers\n",
"Since we've our customer Framework Estimator knows the models are to be served using TensorFlowModel, deploying these endpoints is just a trivial call to the `estimator.deploy()` method"
"## Deploy Models to managed TF Containers\n",
"Since we've serialized the models as TensorFlow SavedModel format, deploying these models as endpoints is just a trivial call to the `estimator.deploy()` method"
]
},
{
@@ -331,13 +305,20 @@
"stax_predictor.delete_endpoint()\n",
"trax_predictor.delete_endpoint()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"display_name": "conda_tensorflow2_p38",
"language": "python",
"name": "python3"
"name": "conda_tensorflow2_p38"
},
"language_info": {
"codemirror_mode": {
@@ -349,9 +330,9 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.5"
"version": "3.8.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
}
advanced_functionality/jax_bring_your_own/training_scripts/requirements.txt
+5
View File
@@ -0,0 +1,5 @@
--find-links https://pypi.org/simple/
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jaxlib==0.3.22+cuda11.cudnn82
jax==0.3.22
trax==1.4.1
advanced_functionality/jax_bring_your_own/training_scripts/train_jax.py
+15 -14
View File
@@ -13,6 +13,13 @@
"""
Train JAX model and serialize as TF SavedModel
"""
# Do not allow TF to take all GPU memory
import os
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
import argparse
import functools
import time
@@ -44,9 +51,7 @@ class PureJaxMNIST:
@staticmethod
def predict(params, inputs, with_classifier=True):
x = inputs.reshape(
(inputs.shape[0], np.prod(inputs.shape[1:]))
) # flatten to f32[B, 784]
x = inputs.reshape((inputs.shape[0], np.prod(inputs.shape[1:]))) # flatten to f32[B, 784]
for w, b in params[:-1]:
x = jnp.dot(x, w) + b
x = jnp.tanh(x)
@@ -70,18 +75,13 @@ class PureJaxMNIST:
predicted_class = jnp.argmax(predict(params, inputs), axis=1)
return jnp.mean(predicted_class == target_class)
batched = [
_per_batch(inputs, labels) for inputs, labels in tfds.as_numpy(dataset)
]
batched = [_per_batch(inputs, labels) for inputs, labels in tfds.as_numpy(dataset)]
return jnp.mean(jnp.stack(batched))
@staticmethod
def update(params, step_size, inputs, labels):
grads = jax.grad(PureJaxMNIST.loss)(params, inputs, labels)
return [
(w - step_size * dw, b - step_size * db)
for (w, b), (dw, db) in zip(params, grads)
]
return [(w - step_size * dw, b - step_size * db) for (w, b), (dw, db) in zip(params, grads)]
@staticmethod
def train(train_ds, test_ds, num_epochs, step_size, with_classifier=True):
@@ -117,14 +117,14 @@ def save_model_tf(prediction_function, params_to_save):
tf_graph = tf.function(
lambda inputs: tf_fun(param_vars, inputs),
autograph=False,
experimental_compile=True,
jit_compile=False,
)
# This signature is needed for TensorFlow Serving use.
signatures = {}
signatures[
tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
] = tf_graph.get_concrete_function(tf.TensorSpec((1, 28, 28, 1), tf.float32))
signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = tf_graph.get_concrete_function(
tf.TensorSpec((1, 28, 28, 1), tf.float32)
)
wrapper = _ReusableSavedModelWrapper(tf_graph, param_vars)
model_dir = "/opt/ml/model/1"
@@ -157,6 +157,7 @@ class _ReusableSavedModelWrapper(tf.train.Checkpoint):
def _parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model_dir", type=str)
parser.add_argument("--num_epochs", type=int, default=3)
parser.add_argument("--batch_size", type=int, default=16)
parser.add_argument("--learning_rate", type=float, default=0.001)
advanced_functionality/jax_bring_your_own/training_scripts/train_stax.py
+16 -14
View File
@@ -14,15 +14,20 @@
Train JAX model using purely functional code and serialize as TF SavedModel
"""
# Do not allow TF to take all GPU memory
import os
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
import argparse
import time
from itertools import count
from jax import random, grad, jit, numpy as jnp
from jax.experimental import optimizers
from jax.experimental import stax
from jax.experimental.stax import Conv, Dense, Relu, LogSoftmax, Flatten
from jax.example_libraries import optimizers
from jax.example_libraries.stax import Conv, Dense, Relu, LogSoftmax, Flatten, serial
import tensorflow as tf
import tensorflow_datasets as tfds
from jax.experimental import jax2tf
@@ -56,7 +61,7 @@ def init_nn():
LogSoftmax,
]
return stax.serial(*layers)
return serial(*layers)
def get_acc_loss_and_update_fns(predict_fn, opt_update, get_params):
@@ -101,9 +106,7 @@ def train(train_ds, test_ds, num_epochs, step_size):
opt_state = opt_init(init_params)
itercount = count()
accuracy, loss, update = get_acc_loss_and_update_fns(
predict_fn, opt_update, get_params
)
accuracy, loss, update = get_acc_loss_and_update_fns(predict_fn, opt_update, get_params)
for epoch in range(num_epochs):
start_time = time.time()
@@ -130,14 +133,14 @@ def save_model_tf(prediction_function, params_to_save):
tf_graph = tf.function(
lambda inputs: tf_fun(param_vars, inputs),
autograph=False,
experimental_compile=True,
jit_compile=False,
)
# This signature is needed for TensorFlow Serving use.
signatures = {}
signatures[
tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
] = tf_graph.get_concrete_function(tf.TensorSpec((1, 28, 28, 1), tf.float32))
signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = tf_graph.get_concrete_function(
tf.TensorSpec((1, 28, 28, 1), tf.float32)
)
wrapper = _ReusableSavedModelWrapper(tf_graph, param_vars)
model_dir = "/opt/ml/model/1"
@@ -170,6 +173,7 @@ class _ReusableSavedModelWrapper(tf.train.Checkpoint):
def _parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model_dir", type=str)
parser.add_argument("--num_epochs", type=int, default=3)
parser.add_argument("--batch_size", type=int, default=16)
parser.add_argument("--learning_rate", type=float, default=0.001)
@@ -183,9 +187,7 @@ if __name__ == "__main__":
train_ds = load_fashion_mnist(tfds.Split.TRAIN, batch_size=args.batch_size)
test_ds = load_fashion_mnist(tfds.Split.TEST, batch_size=args.batch_size)
predict_fn, final_params = train(
train_ds, test_ds, args.num_epochs, args.learning_rate
)
predict_fn, final_params = train(train_ds, test_ds, args.num_epochs, args.learning_rate)
print("finished training")
advanced_functionality/jax_bring_your_own/training_scripts/train_trax.py
+7 -1
View File
@@ -13,8 +13,13 @@
"""
Train Trax model and serialize as TF SavedModel
"""
import argparse
# Do not allow TF to take all GPU memory
import os
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
import argparse
import tensorflow as tf
import trax
from trax import layers as tl
@@ -57,6 +62,7 @@ def save_model_tf(model_to_save):
def _parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--model_dir", type=str)
parser.add_argument("--train_steps", type=int, default=50)
parser.add_argument("--learning_rate", type=float, default=0.001)