在之前的文章中,我们了解了如何转换 Pytorch 模型并使用 OpenVINO 运行,以及OpenVINO 工具包中有哪些深度学习模型优化工具。今天,让我们看看如何使用 TensorFlow 训练的模型并部署到 OpenVINO 推理引擎中运行。
概述
- 设置环境
- 准备 TensorFlow 模型
- 将模型转换为中间表示格式
- 在 OpenVINO 中运行模型推理
- 结论
设置环境
首先,我们需要准备一个 Python 环境:Python 3.5 或更高版本(根据系统要求)和 virtualenv 是我们需要的。
python3 -m venv ~/venv/tf_openvino source ~/venv/tf_openvino/bin/activate
然后让我们安装所需的软件包。
pip3 install --upgrade pip setuptools pip3 install -r requirements.txt
这里 requirements.txt 包含以下软件包。
numpy tqdm tensorflow-cpu==1.15 argparse scipy imageio moviepy
我们需要做的第二件事是使用官方安装说明(在我的例子中是 Linux)安装最新的 OpenVINO 发行版。此外,请记住设置所需的环境变量。
source /opt/intel/openvino_2021/bin/setupvars.sh
将其操作添加到 shell 初始化脚本或虚拟环境激活脚本中可能很有用,以便默认情况下触发它。
准备 TensorFlow 模型
您是否曾经希望一位著名的艺术家用独特的美学视野和风格为您的心爱的猫咪画一幅画?深度神经网络时代的工具发展使得最初在艺术风格的神经算法中提出的神经风格迁移算法成为可能。它的想法是利用神经网络提取图像内容并将其与所使用的表示(风格)分离。然后,我们可以使用提取的风格将其与任意图像结合,并获得令人印象深刻的人工图像结果,这些图像由原始图像的内容和所需的艺术风格组成。
图像风格迁移示例
在上面的示例中,您可以看到如何将弗朗西斯·皮卡比亚的《Udnie》绘画的艺术风格转移到参考猫咪图像(由Dương Nhân拍摄)。显然,要执行此类技巧,我们需要一个经过训练的 TensorFlow 模型。我们可以从头开始准备和训练一个模型,但让我们使用一个预训练的模型。我建议使用来自此仓库的快速风格迁移模型。它非常适合我们的目的。
引用的存储库提供了多个模型检查点,每个检查点都针对不同的艺术风格图像进行训练,可以在这里下载。通常,TensorFlow 检查点包含模型权重和训练期间使用的计算图元数据。为了将模型部署到生产环境中并将其用于推理,我们需要获取冻结图。让我们克隆仓库并使用一个简单的 Python 脚本来执行此操作。
def main(checkpoint_path, input_shape, out_graph_name): # Init graph and session to be used g = tf.Graph() soft_config = tf.compat.v1.ConfigProto(allow_soft_placement=True) with g.as_default(), g.device('/cpu'), tf.compat.v1.Session(config=soft_config) as sess: # Placeholder variable for graph input img_placeholder = tf.compat.v1.placeholder(tf.float32, shape=input_shape, name='img_placeholder') # The model from the repo transform.net(img_placeholder) # Restore model from checkpoint saver = tf.compat.v1.train.Saver() saver.restore(sess, checkpoint_path) # Freeze graph from the session. # "add_37" is the actual last operation of graph frozen = tf.compat.v1.graph_util.convert_variables_to_constants(sess, sess.graph_def, ["add_37"]) # Write frozen graph to a file graph_io.write_graph(frozen, './', out_graph_name, as_text=False) print(f'Frozen graph {out_graph_name} is saved!')
脚本的完整代码可以在这里下载。
只需将脚本放在模型仓库的根文件夹中,并使用以下参数执行它。
python get_frozen_graph.py --checkpoint models/wave.ckpt
在上面的代码片段中,我们执行了所需的 TensorFlow 图预初始化,并使用tf.Saver()从检查点恢复模型数据。然后,我们将转换后的冻结图写入文件。
请注意,冻结图转换需要输出图操作的名称,在本例中为“add_37”。对于任意模型,您可以通过简单地打印网络节点或在 Netron 等工具中探索模型来找到输出操作。
将模型转换为中间表示格式
好的,现在我们有了冻结图,接下来是什么?对于PyTorch 模型,要在 OpenVINO 推理引擎中运行推理,我们必须将模型转换为中间表示 (IR) 格式。幸运的是,OpenVINO 模型优化器内置了对 TensorFlow 模型转换的支持。您可以在此OpenVINO 页面上查看当前支持的 TensorFlow 操作集。
在开始之前,我们应该配置模型优化器(如果我们在 OpenVINO 安装期间跳过了此步骤)。您可以在此处找到配置说明。
由于我们准备好的冻结图已经包含了所有关于模型所需的信息,例如所需的输入形状和输出节点名称,因此我们可以使用默认参数运行模型优化器。
mo_tf.py --input_model inference_graph.pb
如果转换成功,我们将获得包含结果 IR 模型描述的 `inference_graph.xml` 和包含模型权重数据的 `inference_graph.bin`。如果转换不成功,最好从检查目标框架(在本例中为 TensorFlow)的支持层开始。可能需要稍微更改模型架构或为 OpenVINO 实现自定义层。
在 OpenVINO 中运行模型推理
让我们准备一个简单的 Python 脚本,其中包含 OpenVINO 推理引擎初始化、IR 模型加载和提供的图像推理。此脚本的完整版本也在这里提供。
首先,我们导入所需的软件包并定义一个用于参数解析的函数。
import os import cv2 import argparse import time import numpy as np from openvino.inference_engine import IECore from tqdm import tqdm IMG_EXT = ('.png', '.jpg', '.jpeg', '.JPG', '.JPEG') def parse_args(): """Parses arguments.""" parser = argparse.ArgumentParser(description='OpenVINO inference script') parser.add_argument('-i', '--input', type=str, default='', help='Directory to load input images, path to a video or ' 'skip to get stream from the camera (default).') parser.add_argument('-m', '--model', type=str, default='./models/inference_graph.xml', help='Path to IR model') return parser.parse_args()
下一个函数用于 OpenVINO 推理初始化和中间表示模型加载。我们获取 .xml 和 .bin 模型文件的路径,读取它们并将模型加载到推理引擎中。之后,我们可以将其用于推理。
def load_to_IE(model): # Getting the *.bin file location model_bin = model[:-3] + "bin" # Loading the Inference Engine API ie = IECore() # Loading IR files net = ie.read_network(model=model, weights=model_bin) input_shape = net.inputs["img_placeholder"].shape # Loading the network to the inference engine exec_net = ie.load_network(network=net, device_name="CPU") print("IR successfully loaded into Inference Engine.") return exec_net, input_shape
OpenVINO 允许我们使用两种推理请求:同步 - 每个后续请求仅在之前的请求完成后执行,以及异步 - 几个请求由几个执行器同时执行。在此示例中,我们使用同步推理请求接口。
def sync_inference(exec_net, image): input_blob = next(iter(exec_net.inputs)) return exec_net.infer({input_blob: image})
在主函数中,我们首先准备将用于风格迁移的图像列表,然后初始化推理引擎。
def main(args): if os.path.isdir(args.input): # Create a list of test images image_filenames = [os.path.join(args.input, f) for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f)) and f.endswith(IMG_EXT)] image_filenames.sort() else: image_filenames = [args.input] exec_net, net_input_shape = load_to_IE(args.model) # We need dynamically generated key for fetching output tensor output_key = list(exec_net.outputs.keys())[0]
对于每个输入图像,我们将它调整为网络的输入大小,并使用 cv2.dnn.blobFromImage 将其转换为所需的张量格式。然后,最有趣的部分 - 推理请求,然后是简单的后处理。
for image_num in tqdm(range(len(image_filenames))): image = cv2.imread(image_filenames[image_num]) image = cv2.resize(image, (net_input_shape[3], net_input_shape[2])) X = cv2.dnn.blobFromImage(image, swapRB=True) out = sync_inference(exec_net, image=X) result_image = np.squeeze(np.clip(out[output_key], 0, 255).astype(np.uint8), axis=0).transpose((1, 2, 0))
最后,我们显示结果图像。
cv2.imshow("Out", cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)) cv2.waitKey(0) cv2.destroyWindow("Out") if __name__ == "__main__": main(parse_args())
让我们运行脚本,指定要应用风格迁移的图像的输入目录和转换后的 IR 模型。
python openvino_inference.py -i ./images/ -m ./models/inference_graph.xml
带有转移风格的结果图像将在弹出窗口中显示。
性能比较
作为最后一步,让我们比较原始 TensorFlow 和转换为 OpenVINO 推理管道的推理时间。为此,我在推理调用周围添加了一个时间戳测量:在 TensorFlow 的情况下为会话运行调用,在 OpenVINO 的情况下为同步推理请求。在 100 张相同大小(720x1024x3)的图像上平均,每张图像的性能结果非常相似(以秒为单位,CPU:Intel Core i7-8700,12 个线程,4.60 GHz)
OpenVINO_CPU Inference time: Mean: 0.797 Min: 0.705 Max: 1.076 TensorFlow_Cpu Inference time: Mean: 0.816 Min: 0.792 Max: 0.892
但我们可以看到 OpenVINO 版本的速度略快。
结论
在这篇简短的文章中,我们了解了如何轻松地转换 TensorFlow 模型并在 OpenVINO 推理引擎环境中运行它。此外,我们还了解了神经风格迁移算法的应用,并尝试将其应用于任意图像。我希望这是一次有趣且有益的旅程。