手机网站制作代码,微信微商城怎么做,asp.net程序做的网站安全吗,关于建设设计院公司网站的建议在上篇文章中#xff0c;我们详细分享了 JAX 这一新兴的机器学习模型的发展和优势#xff0c;本文我们将通过 Amazon SageMaker 示例展示如何部署并使用 JAX。JAX 的工作机制JAX 的完整工作机制可以用下面这幅图详细解释:图片来源#xff1a;“Intro to JAX” video on YouT…在上篇文章中我们详细分享了 JAX 这一新兴的机器学习模型的发展和优势本文我们将通过 Amazon SageMaker 示例展示如何部署并使用 JAX。JAX 的工作机制JAX 的完整工作机制可以用下面这幅图详细解释:图片来源“Intro to JAX” video on YouTube by Jake VanderPlas, Tech leader from JAX team在图片左侧是开发者自己编写的 Python 代码JAX 会追踪并变换成 JAX IR 的中间表示并按照 Python 代码通过 jax.jit 将其编译成 HLO (High Level Optimized) 代码代表高级的优化代码提供给 XLA 进行读取。XLA 在获取编译的 HLO 代码之后会分配到对应的 CPU、GPU、TPU 或者 ASIC。对于开发者来说只需完成您的 Python 代码即可实现这一流程。开发者可以将 JAX 转换视为首先对 Python 函数进行跟踪专门化然后将其转换为一个小而行为良好的中间形式然后使用特定于转换的解释规则进行解释。为什么 JAX 可以在如此小的软件包中提供如此强大的功能呢首先它从熟悉且灵活的编程接口使用 NumPy 的 Python开始并且使用实际的 Python 解释器来完成大部分繁重的工作其次它将计算的本质提炼成一个静态具有高阶功能的类型表达式语言即 Jaxpr 语言。JAX 应用场景自 2019 年 JAX 出现之后使用它的开发者逐年增多。在 2022 年更是达到了非常火热的状态甚至有人认为它有可能会取代其他的机器学习框架。支持 JAX 生态的应用场景包括深度学习 (Deep Learning)JAX 在深度学习场景下应用很广泛很多团队基于 JAX 开发了更加高级的 API 支持不同的场景方便开发者使用。科学模拟 (Scientific Simulation)JAX 的出现不仅仅是针对于深度学习其实也拥有很多其他的使命如科学模拟。机器人与控制系统 (Robotics and Control Systems)概率编程 (Probabilistic Programming)训练和部署深度学习模型我们用下面这个具体例子展示使用 JAX 来和 Amazon SageMaker 训练和部署深度学习模型会用到 Amazon SageMaker 的 BYOC 这种模式。如上图所示在这个 Amazon SageMaker 的示例中提供了 JAX 的代码示例https://sagemaker-examples.readthedocs.io/en/latest/advanced_functionality/jax_bring_your_own/train_deploy_jax.html在 Amazon SageMaker 上基于 JAX 的框架可使用自定义的容器来训练神经网络。如图的 Amazon SageMaker Examples 提供的 JAX 示例中我们使用自定义容器在 SageMaker 上 基于 JAX 框架或库训练神经网络。这在单个容器上是可能的因为我们使用了 sagemaker-training-toolkit它允许你在自己的自定义容器中使用脚本模式。自定义容器可以使用内置的 SageMaker 训练作业功能如竞价训练和超参数调整。训练模型后您可以将经过训练的模型部署到托管端点。如前所述SageMaker 具有推理容器这些容器已针对亚马逊云科技的硬件和常用深度学习框架进行了优化。其中一项优化是针对 TensorFlow 框架的优化。由于 JAX 支持将模型导出为 TensorFlow SavedModel 格式因此我们使用该功能来展示如何在优化的 SageMaker TensorFlow 推理端点上部署经过训练的模型。整个训练和部署主要分为以下五个步骤创建 Docker 镜像并将其推送到 Amazon ECR。使用 SageMaker 开发工具包传教自定义框架估算器以便将模型输出归类为 TensorFlowModel。代码仓库中有训练估算器的脚本。使用 GPU 上的 SageMaker 训练作业来训练每个模型。将模型部署到完全托管的终端节点。下面我们来看看详细步骤创建 Docker 镜像并将其推送到 Amazon ECR。*创建使用 JAX 训练模型容器的 DockerfileDocker 映像是在 NVIDIA 提供的支持 CUDA 的容器之上构建的。为了确保作为 JAX 中功能基础的 jaxlibpackage 支持 CUDA请从 jax_releases 存储库中下载 jaxlib 软件包。AX releaseshttps://storage.googleapis.com/jax-releases/jax_releases.html这里需要注意的是为了确保作为 JAX 中的功能基础的 JAX library package 能够支持 cuda建议在去做这个创建自定义容器时去看一下目前 JAX release 这个存储库中它下载的这个 JAX library 包的版本号或者相关注意事项等等。2、使用 SageMaker 开发工具包创建自定义框架估算器以便将模型输出归类为 TensorFlowModel。创建基本 SageMaker 框架估算器的子类将估算器的模型类型指定为 TensorFlow 模型。为此我们指定了一个自定义 create_model 方法该方法使用现有的 TensorFlowModel 类来启动推理容器。3、通过代码仓库训练估算器的脚本。您可以通过传统的 SageMaker Python SDK 工作流通过模型执行训练、部署和运行推理。我们确保导入并初始化自定义框架估算器的代码片段中定义的 JaxEstimator然后运行标准的 .fit () 和 .deploy () 调用。对于 JAX 可以调用 jax2tf 函数来执行相同的操作。代码在存储库中可用。设置正确的路径 /opt/ml/model/1 非常重要这是 SageMaker wrapper封装器 假定模型已存储的地方。、前面提到的 JAX 和 TF 的互操作性目前 JAX 是通过 JAX to TF 这样的一个软件包来为 JAX 和 TF 的互操作性提供支持那 jax2tf.convert 是用于在 TensorFlow 的上下文中使用 JAX 函数那 jax2tf.call_tf 是用于在 JAX 的上下文中使用的 TensorFlow 函数互操作来完成的。4、使用 GPU 上的 SageMaker 训练作业来训练每个模型。将模型部署到完全托管的终端节点。vanilla_jax_predictor vanilla_jax_estimator.deploy(initial_instance_count1, instance_typeml.m4.xlarge
)import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt(x_train, y_train), (x_test, y_test) tf.keras.datasets.fashion_mnist.load_data()def test_image(predictor, test_images, test_labels, image_number):np_img np.expand_dims(np.expand_dims(test_images[image_number], axis
-1
), axis
0
)result predictor.predict(np_img)pred_y np.argmax(result[predictions])print(True Label:, test_labels[image_number])print(Predicted Label:, pred_y)plt.imshow(test_images[image_number]*部署和准备输入的测试图像*进行推理有关在 Amazon SageMaker 上使用 JAX 训练和部署深度学习模型的详细过程和代码请参考亚马逊云科技官方博客。如图所示上面的两张图是一个部署模型的例子下面的图是进行推理的例子。由于我们的 Framework Estimator 知道模型将使用 TensorFlowModel 提供服务因此部署这些端点只是对 estimator.deploy () 方法做调用即可。参考资料Training and Deploying ML Models using JAX on SageMakerTrain and deploy deep learning models using JAX with Amazon SageMakerAX core from scratchBuilding JAX from sourceJAX 是一种越来越流行的库它支持原生 Python 或 NumPy 函数的可组合函数转换可用于高性能数值计算和机器学习研究。JAX 提供了编写 NumPy 程序的能力这些程序可以使用 GPU/TPU 自动差分和加速从而形成了更灵活的框架来支持现代深度学习架构。在这两篇文章中我们讨论了有关 JAX 的一些主题希望对您用使用 JAX 这一框架进行深度学习研究有所帮助。往期推荐机器学习洞察 | JAX机器学习领域的“新面孔”机器学习洞察 | 降本增效无服务器推理是怎么做到的机器学习洞察 | 分布式训练让机器学习更加快速准确
