如何在Jetson Nano上运行TensorFlow对象检测模型

  • 0 replies
  • 2010 views
*

sisiy

  • *****
  • 203
    • 查看个人资料
如何在Jetson Nano上运行TensorFlow对象检测模型
« 于: 五月 05, 2019, 02:51:06 pm »
这一次,您将了解如何在JEtson NANO 上运行一个Tensorflow对象检测模型。它可以是Tensorflow detection model zoo中的一个预训练模型(https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md),可以检测日常对象,比如人/车/狗,也可以是一个自定义训练的对象检测模型,可以检测你自定义的对象。

本次教学内容是转换在coco数据集上训练好的,用于通用目标识别的SSD MobileNet V1模型。
下面的内容将分解/分析如何完成这个转换过程。
1.下载预训练的模型检查点,构建TensorFlow检测图,然后使用TensorRT创建推理图。
2.在Jetson Nano上加载tensorrt推理图并进行预测。

这两个步骤将在两个单独的Jupyter笔记本中处理,第一个运行在开发机器上,第二个运行在Jetson Nano上。
在继续之前,请确保您已经安装了Jetson Nano并安装了Tensorflow。

步骤1:创建TensorRT模型
在你的开发机器上用 Tensorflow nightly builds来构建环境,这个默认包括TF-TRT 。
在notebook中,您将首先安装Tensorflow对象检测API并设置相关路径。它的官方安装文档(https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/installation.md)对初学者来说可能有点吓人,但是你也可以通过运行一个notebook cell来完成。
程序代码: [选择]
%cd /content
!git clone --quiet https://github.com/tensorflow/models.git

!apt-get install -qq protobuf-compiler python-pil python-lxml python-tk

!pip install -q Cython contextlib2 pillow lxml matplotlib

!pip install -q pycocotools

%cd /content/models/research
!protoc object_detection/protos/*.proto --python_out=.

import os
import sys
os.environ['PYTHONPATH'] += ':/content/models/research/:/content/models/research/slim/'
sys.path.append("/content/models/research/slim/")

!python object_detection/builders/model_builder_test.py

接下来,您将从预先训练好的ssd_mobilenet_v1_coco检查点下载并构建一个检测图,或者从notebook中提供的列表中选择另一个。
程序代码: [选择]
config_path, checkpoint_path = download_detection_model(MODEL, 'data')

frozen_graph, input_names, output_names = build_detection_graph(
    config=config_path,
    checkpoint=checkpoint_path,
    score_threshold=0.3,
    iou_threshold=0.5,
    batch_size=1
)

最初,默认的Tensorflow对象检测模型采用可变批量大小,现在固定为1,因为Jetson Nano是一个资源受限的设备。在build_detection_graph调用中,其他几个更改应用于Tensorflow图,
1.将得分阈值设置为0.3,模型将剔除任何低于该阈值的预测结果。
2.IoU(交集超过联合)阈值设置为0.5,这样任何检测到的具有相同类重叠的对象都将被删除。您可以在这里(https://www.dlology.com/blog/gentle-guide-on-how-yolo-object-localization-works-with-keras-part-2/)阅读更多关于IoU(union上的交集)和非最大抑制的信息。
3.应用修改冻结对象检测图,以提高速度和减少内存消耗。

接下来,我们创建一个与图像分类模型类似的TensorRT推理图。


程序代码: [选择]
import tensorflow.contrib.tensorrt as trt

trt_graph = trt.create_inference_graph(
    input_graph_def=frozen_graph,
    outputs=output_names,
    max_batch_size=1,
    max_workspace_size_bytes=1 << 25,
    precision_mode='FP16',
    minimum_segment_size=50
)

有了TensorRT推理图之后,可以将其保存为pb文件,并根据需要从Colab或本地机器下载到Jetson Nano中。
程序代码: [选择]
with open('./data/trt_graph.pb', 'wb') as f:
    f.write(trt_graph.SerializeToString())

# Download the tensorRT graph .pb file from colab to your local machine.
from google.colab import files

files.download('./data/trt_graph.pb')

步骤2:加载TensorRT图并进行预测
 
在你的Jetson Nano上,用命令jupyter notebook -ip = 0.0.0.0启动一个Jupyter Notebook,你已经将下载的图形文件保存到./model/trt_graph.pb。 以下代码将加载TensorRT图并使其为推理做好准备。

程序代码: [选择]
import tensorflow as tf

def get_frozen_graph(graph_file):
    """Read Frozen Graph file from disk."""
    with tf.gfile.FastGFile(graph_file, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    return graph_def

# The TensorRT inference graph file downloaded from Colab or your local machine.
pb_fname = "./model/trt_graph.pb"
trt_graph = get_frozen_graph(pb_fname)

input_names = ['image_tensor']

# Create session and load graph
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
tf_sess = tf.Session(config=tf_config)
tf.import_graph_def(trt_graph, name='')

tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ':0')
tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')
tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')
tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')
tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')

现在,我们可以使用图像进行预测,并查看模型是否正确。 请注意,我们将图像的大小调整为300 x 300,但是,您可以尝试其他尺寸或仅保持大小不变,因为图形可以处理可变大小的输入。 但请记住,因为与台式机相比,Jetson中的内存非常小,因此很难拍摄大型图像。
程序代码: [选择]
import cv2
IMAGE_PATH = "./data/dogs.jpg"
image = cv2.imread(IMAGE_PATH)
image = cv2.resize(image, (300, 300))

scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={
    tf_input: image[None, ...]
})
boxes = boxes[0]  # index by 0 to remove batch dimension
scores = scores[0]
classes = classes[0]
num_detections = int(num_detections[0])

如果您以前玩过Tensorflow对象检测API,那么这些输出应该很熟悉。
在这里,结果可能仍然包含不同类标签的重叠预测。例如,同一个对象可以在两个重叠的绑定框中使用两个类来标记。
我们将使用自定义的非最大抑制函数来删除预测分数较低的重叠边界框。
让我们通过绘制包围框和覆盖标签来可视化结果。
下面是在Jetson Nano笔记本上创建覆盖和显示的代码。
程序代码: [选择]
from IPython.display import Image as DisplayImage

# Boxes unit in pixels (image coordinates).
boxes_pixels = []
for i in range(num_detections):
    # scale box to image coordinates
    box = boxes[i] * np.array([image.shape[0],
                               image.shape[1], image.shape[0], image.shape[1]])
    box = np.round(box).astype(int)
    boxes_pixels.append(box)
boxes_pixels = np.array(boxes_pixels)

# Remove overlapping boxes with non-max suppression, return picked indexes.
pick = non_max_suppression(boxes_pixels, scores[:num_detections], 0.5)


for i in pick:
    box = boxes_pixels[i]
    box = np.round(box).astype(int)
    # Draw bounding box.
    image = cv2.rectangle(
        image, (box[1], box[0]), (box[3], box[2]), (0, 255, 0), 2)
    label = "{}:{:.2f}".format(int(classes[i]), scores[i])
    # Draw label (class index and probability).
    draw_label(image, (box[1], box[0]), label)

# Save and display the labeled image.
save_image(image[:, :, ::-1])
DisplayImage(filename="./data/img.png")

在coco标签地图中,第18类表示狗,23表示熊。 坐在那里的两只狗被错误地归类为熊。 也许在可可数据集中有比坐着的狗更多的坐熊。
进行了类似的速度基准测试,Jetson Nano通过SSD MobileNet V1型号和300 x 300输入图像实现了11.54 FPS。

结论和进一步阅读
在本教程中,您学习了如何转换Tensorflow对象检测模型并在Jetson Nano上运行推理。
查看更新的GitHub repo(https://github.com/Tony607/tf_jetson_nano)以获取源代码。
如果您对结果不满意,还有其他预先训练好的模型供您查看,我建议您从SSD MobileNet V2(ssd_mobilenet_v2_coco)开始,或者如果您喜欢冒险,请尝试ssd_inception_v2_coco,这可能会突破极限 Jetson Nano的记忆。
您可以在 https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md中找到这些模型,“速度(ms)”指标将为您提供模型复杂性的指导。
考虑使用免费数据中心GPU培训自定义对象检测模型,请查看我之前的教程 https://www.dlology.com/blog/how-to-train-an-object-detection-model-easy-for-free/- 如何轻松免费培训对象检测模型。
« 最后编辑时间: 五月 05, 2019, 03:49:48 pm 作者 sisiy »