在上一篇文章中，我们详细分享了JAX这一新兴的机器学习模型的发展历程和优势。在本文中，我们将展示如何通过AmazonSageMaker示例JAX部署和使用它。JAX的工作机制JAX完整的工作机制可以用下图来详细解释：图片来源：YouTube上的“IntrotoJAX”videobyJakeVanderPlas,TechleaderfromJAXteam图片左边是开发者自己编写的Python代码，JAX会跟踪转化为JAXIR的中间表示，并根据Python代码通过jax.jit编译成HLO（HighLevelOptimized）代码，代表高级优化代码，并提供它到XLA进行阅读。XLA获得编译后的HLO代码后，会分配给对应的CPU、GPU、TPU或ASIC。对于开发人员来说，这个过程可以简单地通过完成你的Python代码来完成。开发人员可以将JAX转换视为首先对Python函数进行跟踪特化，然后将其转换为小型、行为良好的中间形式，然后使用特定于转换的解释规则对其进行解释。为什么JAX可以在如此小的包中提供如此多的功能？首先，它以熟悉且灵活的编程界面（使用NumPy的Python）开始，并使用实际的Python解释器来完成大部分繁重的工作；其次，它将计算的本质提炼成具有高阶函数的静态类型表达式语言，即Jaxpr语言。JAX应用场景自2019年JAX出现以来，使用它的开发者数量逐年增加。它将在2022年达到非常火热的状态，甚至有人认为它可能会取代其他机器学习框架。支持JAX生态的应用场景包括：深度学习：JAX广泛应用于深度学习场景，很多团队基于JAX开发了更高级的API来支持不同的场景，方便开发者使用。科学模拟：JAX的出现不仅仅是为了深度学习，还有很多其他的使命，比如科学模拟。机器人和控制系统概率编程（ProbabilisticProgramming）深度学习模型的训练和部署我们用下面的具体例子来展示使用JAX来训练和部署带有AmazonSageMaker的深度学习模型，AmazonSageMaker的BYOC会用到。模式。如上所示，此AmazonSageMaker示例中提供了JAX的代码示例：https://sagemaker-examples.re...AmazonSageMaker上基于JAX的框架可以使用自定义容器来训练神经网络。在AmazonSageMakerExamples提供的JAX示例中，我们使用自定义容器来训练基于SageMaker上的JAX框架或库的神经网络。这在单个容器上是可能的，因为我们使用sagemaker-training-toolkit，它允许您在自己的自定义容器中使用脚本模式。自定义容器可以使用内置的SageMaker训练作业功能，例如出价训练和超参数调整。训练模型后，您可以将训练后的模型部署到托管端点。如前所述，SageMaker具有针对AmazonCloudTechnology的硬件和通用深度学习框架进行优化的推理容器。其中一项优化是针对TensorFlow框架的。由于JAX支持将模型导出为TensorFlowSavedModel格式，我们使用此功能来展示如何在优化的SageMakerTensorFlow推理端点上部署经过训练的模型。整个训练和部署主要分为以下五个步骤：1.创建Docker镜像并推送到AmazonECR。2.使用SageMakerSDK训练自定义框架估算器将模型输出分类到TensorFlowModel中。3.代码库中有训练估算器的脚本。4.在GPU上使用SageMaker训练作业训练每个模型。5.将模型部署到完全托管的端点。下面我们来看一下详细步骤：1.创建一个Docker镜像并推送到AmazonECR。*用于使用JAX为训练模型创建容器的DockerfileDocker映像构建在NVIDIA提供的支持CUDA的容器之上。为确保作为JAX功能基础的jaxlib包支持CUDA，请从jax_releases存储库下载jaxlib包。AX发布https://storage.googleapis.co...这里需要注意的是：为了保证作为JAX中功能基础的JAX库包能够支持cuda，建议采取一个查看当前JAX版本库中，其下载的JAX库包的版本号或相关注意事项。2.使用SageMakerSDK创建自定义帧估计器，将模型输出分类到TensorFlowModel中。创建基本SageMaker框架估算器的子类，将估算器的模型类型指定为TensorFlow模型。为此，我们指定了一个自定义的create_model方法，该方法使用现有的TensorFlowModel类来启动推理容器。3.通过代码库训练estimator的脚本。您可以通过传统的SageMakerPythonSDK工作流对您的模型执行训练、部署和运行推理。我们确保导入并初始化在自定义框架估算器的代码片段中定义的JaxEstimator，然后运行标准的.fit()和.deploy()调用。对于JAX，可以调用jax2tf函数来执行相同的操作。代码在存储库中可用。设置正确的路径/opt/ml/model/1非常重要，这是SageMaker包装器假定存储模型的位置。,上面提到了JAX和TF的互操作，目前JAX通过JAX转TF等软件包提供了对JAX和TF互操作的支持，而jax2tf.convert是在TensorFlow上下文中使用JAX函数，jax2tf.call_tf是用于与JAX上下文中使用的TensorFlow函数进行互操作。4.在GPU上使用SageMaker训练作业训练每个模型。5.将模型部署到完全托管的端点。vanilla_jax_predictor=vanilla_jax_estimator.deploy(initial_instance_count=1,instance_type="ml.m4.xlarge")将tensorflow导入为tfi将numpy导入为npfrommatplotlib将pyplot导入为plt(x_train,y_train),(x_test,y_test)=tf.keras.datasetsfashion_mnist.load_data()deftest_image(预测器,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("TrueLabel:",test_labels[image_number])print("PredictedLabel:",pred_y)plt.imshow(test_images[image_number]*已部署并准备好输入测试image*forinference在AmazonSageMaker上使用JAX训练和部署深度学习模型的详细过程和代码，请参考亚马逊云技术官方博客。如图所示，以上两张图为部署示例模型，下面的图是做推理的例子。我们的FrameworkEstimator知道模型将使用TensorFlowModel提供服务，部署这些端点只是调用estimator.deploy()方法。参考在SageMaker上使用JAX训练和部署ML模型使用JAX和AmazonSageMakerAX核心从头开始训练和部署深度学习模型从头开始构建JAXJAX是一个越来越受欢迎的库，它支持原生Python或NumPy函数的可组合函数转换，它可用于高性能数值计算和机器学习研究。JAX提供了编写NumPy程序的能力，这些程序可以使用GPU/TPU自动微分和加速，从而产生更灵活的框架来支持现代深度学习架构。在这两篇文章中我们讨论了一些关于JAX的话题，希望对您使用JAX作为深度学习研究的框架有所帮助。之前推荐的机器学习见解|JAX，机器学习领域的“新面孔”Machinelearninginsights|降本增效，Serverless推理是如何实现的？机器学习见解|分布式训练让机器学习更快更准点击查看原文
