准备模型进行编译 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

准备模型进行编译

SageMaker Neo 要求机器学习模型满足特定的输入数据形状。编译所需的输入形状取决于您使用的深度学习框架。模型输入形状格式正确后,请根据以下要求保存模型。保存模型后,压缩模型构件。

SageMaker Neo 需要什么输入数据形状?

在编译模型之前,请确保模型格式正确。Neo 需要已训练模型的预期数据输入(JSON 格式或列表格式)的名称和形状。需要的输入特定于框架。

以下是 SageMaker Neo 需要的输入形状:

使用已训练模型的字典格式,指定预期数据输入的名称和形状(NCHW 格式)。请注意,虽然 Keras 模型构件应以 NHWC (channel-last) 格式上传,但 DataInputConfig 应以 NCHW (channel-first) 格式指定。所需的字典格式如下:

  • 对于一个输入:{'input_1':[1,3,224,224]}

  • 对于两个输入:{'input_1': [1,3,224,224], 'input_2':[1,3,224,224]}

使用已训练模型的字典格式,指定预期数据输入的名称和形状(NCHW 格式)。所需的字典格式如下:

  • 对于一个输入:{'data':[1,3,1024,1024]}

  • 对于两个输入:{'var1': [1,1,28,28], 'var2':[1,1,28,28]}

对于 PyTorch 模型,如果您满足以下两个条件,则无需提供预期数据输入的名称和形状:

  • 您使用 PyTorch 2.0 或更高版本创建了模型定义文件。有关如何创建定义文件的更多信息,请参阅为 SageMaker Neo 保存模型下的PyTorch 部分。

  • 您正在为云实例编译模型。有关 SageMaker Neo 支持的实例类型的更多信息,请参阅 支持的实例类型和框架

如果您满足这些条件,SageMaker Neo 会从您使用 PyTorch 创建的模型定义文件(.pt 或 .pth)中获取输入配置。

否则,您必须执行以下操作:

使用已训练模型的字典格式,指定预期数据输入的名称和形状(NCHW 格式)。或者,您只能使用列表格式指定形状。所需的字典格式如下:

  • 字典格式的一个输入:{'input0':[1,3,224,224]}

  • 列表格式的一个输入:[[1,3,224,224]]

  • 字典格式的两个输入:{'input0':[1,3,224,224], 'input1':[1,3,224,224]}

  • 列表格式的两个输入:[[1,3,224,224], [1,3,224,224]]

使用已训练模型的字典格式,指定预期数据输入的名称和形状(NHWC 格式)。所需的字典格式如下:

  • 对于一个输入:{'input':[1,1024,1024,3]}

  • 对于两个输入:{'data1': [1,28,28,1], 'data2':[1,28,28,1]}

使用已训练模型的字典格式,指定预期数据输入的名称和形状(NHWC 格式)。所需的字典格式如下:

  • 对于一个输入:{'input':[1,224,224,3]}

注意

SageMaker Neo 仅为边缘设备目标支持 TensorFlow Lite。有关受支持的 SageMaker Neo 边缘设备目标的列表,请参阅 SageMaker Neo 设备 页面。有关支持的 SageMaker Neo 云实例目标列表,请参阅 SageMaker Neo 支持的实例类型和框架页面。

不需要输入数据名称和形状。

为 SageMaker Neo 保存模型

以下代码示例展示了如何保存模型以使其与 Neo 兼容。必须将模型打包为压缩 tar 文件 (*.tar.gz)。

Keras 模型需要一个模型定义文件 (.h5)。

为了使其与 SageMaker Neo 兼容,有两种方法可以保存 Keras 模型:

  1. 使用 model.save("<model-name>", save_format="h5") 导出为 .h5 格式。

  2. 导出 SavedModel 后将其冻结。

以下是如何导出 tf.keras 模型作为冻结图表的示例(选项二):

import os import tensorflow as tf from tensorflow.keras.applications.resnet50 import ResNet50 from tensorflow.keras import backend tf.keras.backend.set_learning_phase(0) model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling='avg') model.summary() # Save as a SavedModel export_dir = 'saved_model/' model.save(export_dir, save_format='tf') # Freeze saved model input_node_names = [inp.name.split(":")[0] for inp in model.inputs] output_node_names = [output.name.split(":")[0] for output in model.outputs] print("Input names: ", input_node_names) with tf.Session() as sess: loaded = tf.saved_model.load(sess, export_dir=export_dir, tags=["serve"]) frozen_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), output_node_names) tf.io.write_graph(graph_or_graph_def=frozen_graph, logdir=".", name="frozen_graph.pb", as_text=False) import tarfile tar = tarfile.open("frozen_graph.tar.gz", "w:gz") tar.add("frozen_graph.pb") tar.close()
警告

请勿使用 model.save(<path>, save_format='tf') 将模型导出为 SavedModel 类。这种格式适合训练,但不适合推理。

MXNet 模型必须保存为单个符号文件 *-symbol.json 和单个参数 *.params files

Gluon Models

使用 HybridSequential 类定义神经网络。这将以符号编程的风格运行代码(而不是命令式编程)。

from mxnet import nd, sym from mxnet.gluon import nn def get_net(): net = nn.HybridSequential() # Here we use the class HybridSequential. net.add(nn.Dense(256, activation='relu'), nn.Dense(128, activation='relu'), nn.Dense(2)) net.initialize() return net # Define an input to compute a forward calculation. x = nd.random.normal(shape=(1, 512)) net = get_net() # During the forward calculation, the neural network will automatically infer # the shape of the weight parameters of all the layers based on the shape of # the input. net(x) # hybridize model net.hybridize() net(x) # export model net.export('<model_name>') # this will create model-symbol.json and model-0000.params files import tarfile tar = tarfile.open("<model_name>.tar.gz", "w:gz") for name in ["<model_name>-0000.params", "<model_name>-symbol.json"]: tar.add(name) tar.close()

有关混合模型的更多信息,请参阅 MXNet 混合文档

Gluon Model Zoo (GluonCV)

GluonCV model zoo 模型是预混合的。所以,您可以直接导出它们。

import numpy as np import mxnet as mx import gluoncv as gcv from gluoncv.utils import export_block import tarfile net = gcv.model_zoo.get_model('<model_name>', pretrained=True) # For example, choose <model_name> as resnet18_v1 export_block('<model_name>', net, preprocess=True, layout='HWC') tar = tarfile.open("<model_name>.tar.gz", "w:gz") for name in ["<model_name>-0000.params", "<model_name>-symbol.json"]: tar.add(name) tar.close()
Non Gluon Models

所有非 Gluon 模型使用 *-symbol*.params 文件保存到磁盘。因此,它们已经是 Neo 需要的正确格式。

# Pass the following 3 parameters: sym, args, aux mx.model.save_checkpoint('<model_name>',0,sym,args,aux) # this will create <model_name>-symbol.json and <model_name>-0000.params files import tarfile tar = tarfile.open("<model_name>.tar.gz", "w:gz") for name in ["<model_name>-0000.params", "<model_name>-symbol.json"]: tar.add(name) tar.close()

PyTorch 模型必须保存为定义文件(.pt.pth),输入数据类型为 float32

要保存模型,请使用 torch.jit.trace 方法,再使用 torch.save 方法。此过程将对象保存到磁盘文件中,默认情况下使用 python pickle (pickle_module=pickle) 来保存对象和一些元数据。接下来,将保存的模型转换为压缩的 tar 文件。

import torchvision import torch model = torchvision.models.resnet18(pretrained=True) model.eval() inp = torch.rand(1, 3, 224, 224) model_trace = torch.jit.trace(model, inp) # Save your model. The following code saves it with the .pth file extension model_trace.save('model.pth') # Save as a compressed tar file import tarfile with tarfile.open('model.tar.gz', 'w:gz') as f: f.add('model.pth') f.close()

如果您使用 PyTorch 2.0 或更高版本保存模型,SageMaker Neo 会从定义文件中获取模型的输入配置(其输入的名称和形状)。在这种情况下,编译模型时不需要为 SageMaker 指定数据输入配置。

如果要阻止 SageMaker Neo 获取输入配置,可以将 torch.jit.trace_store_inputs 参数设置为 False。如果进行该设置,编译模型时不需要为 SageMaker 指定数据输入配置。

有关 torch.jit.trace 方法的更多信息,请参阅 PyTorch 文档中的 TORCH.JIT.TRACE

TensorFlow 需要一个 .pb 或者一个 .pbtxt 文件和包含变量的变量目录。对于冷冻模型,只有一个 .pb.pbtxt 文件是必需的。

以下代码示例显示如何使用 tar Linux 命令压缩模型。在终端或 Jupyter 笔记本中运行以下命令(如果您使用的是 Jupyter 笔记本,请在语句开头插入 ! 魔法命令):

# Download SSD_Mobilenet trained model !wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz # unzip the compressed tar file !tar xvf ssd_mobilenet_v2_coco_2018_03_29.tar.gz # Compress the tar file and save it in a directory called 'model.tar.gz' !tar czvf model.tar.gz ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb

本示例中使用的命令标志可完成以下操作:

  • c:创建存档

  • z:使用 gzip 压缩存档

  • v:显示存档进度

  • f:指定存档的文件名

内置估算器由特定于框架的容器或特定于算法的容器完成。使用内置 .fit 方法训练模型时,内置算法和特定于框架的估算器的估算器对象都会以正确的格式保存模型。

例如,您可以使用 sagemaker.TensorFlow 来定义 TensorFlow 估算器:

from sagemaker.tensorflow import TensorFlow estimator = TensorFlow(entry_point='mnist.py', role=role, #param role can be arn of a sagemaker execution role framework_version='1.15.3', py_version='py3', training_steps=1000, evaluation_steps=100, instance_count=2, instance_type='ml.c4.xlarge')

然后用 .fit 内置方法训练模型:

estimator.fit(inputs)

在最后使用内置 compile_model 方法来编译模型之前:

# Specify output path of the compiled model output_path = '/'.join(estimator.output_path.split('/')[:-1]) # Compile model optimized_estimator = estimator.compile_model(target_instance_family='ml_c5', input_shape={'data':[1, 784]}, # Batch size 1, 3 channels, 224x224 Images. output_path=output_path, framework='tensorflow', framework_version='1.15.3')

您也可以利用 SageMaker Python SDK 中的 compile_model 方法,使用 sagemaker.estimator.Estimator 类初始化一个估算器对象以训练和编译内置算法:

import sagemaker from sagemaker.image_uris import retrieve sagemaker_session = sagemaker.Session() aws_region = sagemaker_session.boto_region_name # Specify built-in algorithm training image training_image = retrieve(framework='image-classification', region=aws_region, image_scope='training') training_image = retrieve(framework='image-classification', region=aws_region, image_scope='training') # Create estimator object for training estimator = sagemaker.estimator.Estimator(image_uri=training_image, role=role, #param role can be arn of a sagemaker execution role instance_count=1, instance_type='ml.p3.8xlarge', volume_size = 50, max_run = 360000, input_mode= 'File', output_path=s3_training_output_location, base_job_name='image-classification-training' ) # Setup the input data_channels to be used later for training. train_data = sagemaker.inputs.TrainingInput(s3_training_data_location, content_type='application/x-recordio', s3_data_type='S3Prefix') validation_data = sagemaker.inputs.TrainingInput(s3_validation_data_location, content_type='application/x-recordio', s3_data_type='S3Prefix') data_channels = {'train': train_data, 'validation': validation_data} # Train model estimator.fit(inputs=data_channels, logs=True) # Compile model with Neo optimized_estimator = estimator.compile_model(target_instance_family='ml_c5', input_shape={'data':[1, 3, 224, 224], 'softmax_label':[1]}, output_path=s3_compilation_output_location, framework='mxnet', framework_version='1.7')

有关使用 SageMaker Python SDK 编译模型的更多信息,请参阅 编译模型 (Amazon SageMaker SDK)