MegEngine Lite Python 部署模型快速上手#
Note
Lite 的 Python 包是和 MegEngine 本体绑定在一起的,所以只需要 安装 MegEngine 包 即可使用。
相较于 C++ 推理,该方法可以直接加载导出的模型并执行推理,无需经历复杂的编译环节。
Warning
但是这种方式目前只支持 Windows/MacOS/Linux x86 和 CUDA 版本,不支持 Android Arm.
本文将从获取一个训练好的 shufflenet_v2
模型出发,讲解如何使用 MegEngine Lite 的 Python 接口将其部署到 Linux x86 中环境下运行。
主要分为以下小节:
导出已经训练好的模型#
编写并执行 Inference 代码#
首先创建一个 inference.py
, 在这个文件中将直接调用 MegEngine Lite 的 Python 接口运行 shufflenet_v2.mge
模型,
注意输入数据 input_tensor
是随机生成的,不用在乎计算结果。
from megenginelite import *
import numpy as np
import argparse
def main():
parser = argparse.ArgumentParser(description='test the lite python interface')
parser.add_argument('input', help='the inference model file')
args = parser.parse_args()
network = LiteNetwork()
network.load(args.input)
input_tensor = network.get_io_tensor("data")
# copy input data to input_tensor of the network
input_data = np.random.uniform(0, 1, (1, 3, 224, 224)).astype("float32")
input_tensor.set_data_by_copy(input_data)
# forward the model
network.forward()
network.wait()
output_names = network.get_all_output_name()
output_tensor = network.get_io_tensor(output_names[0])
output_data = output_tensor.to_numpy()
print('shufflenet output max={}, sum={}'.format(output_data.max(), output_data.sum()))
if __name__ == '__main__':
main()
上面代码主要完成了几个步骤,包括:
创建默认配置的 Network;
载入模型,MegEngine Lite 将读取并解析模型文件,并创建计算图;
通过输入 Tensor 的名字获取模型的输入 Tensor, 并设置随机数作为输入数据;
执行 Inference 逻辑;
获取模型输出 Tensor, 并处理输出数据。
这样这个调用 MegEngine Lite 的 Python 接口的 demo 就完成了。
使用 CUDA 进行推理#
下面将通过 CUDA 运行 shufflenet.mge 来展示如何使用 TensorBatchCollector 来攒 batch,攒好之后传递到 network 的输入 tensor 中进行推理。
from megenginelite import *
import numpy as np
import argparse
def main():
parser = argparse.ArgumentParser(description='test the lite python interface')
parser.add_argument('input', help='the inference model file')
args = parser.parse_args()
# construct LiteOption
net_config = LiteConfig(device_type=LiteDeviceType.LITE_CUDA)
network = LiteNetwork(config=net_config)
network.load(args.input)
input_tensor = network.get_io_tensor("data")
# copy input data to input_tensor of the network
input_data = np.random.uniform(0, 1, (1, 3, 224, 224)).astype("float32")
input_tensor.set_data_by_copy(input_data)
# forward the model
network.forward()
network.wait()
output_names = network.get_all_output_name()
output_tensor = network.get_io_tensor(output_names[0])
output_data = output_tensor.to_numpy()
print('shufflenet output max={}, sum={}'.format(output_data.max(), output_data.sum()))
if __name__ == '__main__':
main()
上面示例主要演示在 CUDA 设备上进行 Inference 的接口调用过程。
使用 TensorBatchCollector 辅助完成 CUDA 推理#
下面将通过 CUDA 运行 shufflenet.mge 来展示如何使用 TensorBatchCollector 来攒 batch,攒好之后传递到 network 的输入 tensor 中进行推理。
from megenginelite import *
import numpy as np
import os
def test_network():
model_path = "shufflenet.mge"
batch = 4
# construct LiteOption
net_config = LiteConfig(device_type=LiteDeviceType.LITE_CUDA)
# constuct LiteIO, is_host=False means the input tensor will use device memory
# set the input tensor "data" memory is not from host, but from device
ios = LiteNetworkIO(inputs=[["data", is_host=False]])
network = LiteNetwork(config=net_config, io=ios)
network.load(model_path)
dev_input_tensor = network.get_io_tensor("data")
# read input to input_data
input_layout = dev_input_tensor.layout
shape = list(input_layout.shapes)[0 : input_layout.ndim]
arr = np.ones(shape[1:], "float32")
shape[0] = batch
print(shape)
batch_tensor = TensorBatchCollector(
shape, dtype=LiteDataType.LITE_FLOAT, device_type=LiteDeviceType.LITE_CUDA
)
for time in range(3):
batch_tensor.free(range(batch))
for i in range(batch):
batch = batch_tensor.collect(arr)
print("collect batch id = {}".format(batch))
arr += 1
# set device input data to input_tensor of the network without copy
dev_input_tensor.share_memory_with(batch_tensor.get())
network.forward()
network.wait()
output_names = network.get_all_output_name()
output_tensor = network.get_io_tensor(output_names[0])
output_data = output_tensor.to_numpy()
print('shufflenet output shape={}, max={}, sum={}'.format(output_data.shape, output_data.max(), output_data.sum()))
test_network()
上面示例主要做了以下事情:
通过 LiteConfig 和 LiteIO 来创建一个运行在 CUDA 上的 Network,并配置该 Network 中输入名字为 “data” 的 Tensor 在 CUDA 上,这样用户可以直接将 CUDA device 上的内存 share 给它。
通过该 Network 加载 shufflenet 模型,并获取名字为 “data” 的输入 Tensor,以及它的 layout 信息。
通过输入 tensor 的 layout 信息和 batch 信息,将创建一个在 CUDA 上的 TensorBatchCollector,并循环攒了 4 个 batch。
然后将 TensorBatchCollector 中的 tensor 和 Network 的输入 tensor 通过 share_memory_with 进行内存 share。
执行推理,获取输出数据
Note
上面通过 share_memory_with 进行内存共享,将不会产生多余的数据 copy,其中 TensorBatchCollector 的使用请参考 Utils API。