如何对 MXNet 和 Python 2.7 中的多个图像同时对图像分类模型进行推理

Question

我正在使用 Python 2.7、MXNet V1.3.0 ML 框架在 ONNX 格式的图像分类模型（带有 opset 7 的 V1.2.1）上运行推理，我一次将图像提供给推理器。我需要做什么才能异步运行多个图像的推理但还要等待所有图像完成？

我从 30 FPS 的视频中提取帧作为 .jpeg 图像。例如，当我对长度为 20 秒的视频运行该过程时，它会生成 600 个 .jpeg 图像。现在，我遍历这些图像的列表，并将它们中的每一个的相对路径传递给以下函数，然后从目标图像推断。

def infer(self, target_image_path):
        target_image_path = self.__output_directory + '/' + target_image_path

        image_data = self.__get_image_data(target_image_path)  # Get pixel data

        '''Define the model's input'''
        model_metadata = onnx_mxnet.get_model_metadata(self.__model)
        data_names = [inputs[0]
                      for inputs in model_metadata.get('input_tensor_data')]
        Batch = namedtuple('Batch', 'data')

        ctx = mx.eia()  # Set the context to elastic inference

        '''Load the model'''
        sym, arg, aux = onnx_mxnet.import_model(self.__model)
        mod = mx.mod.Module(symbol=sym, data_names=data_names,
                            context=ctx, label_names=None)
        mod.bind(data_shapes=[(data_names[0], image_data.shape)],
                 label_shapes=None, for_training=False)

        mod.set_params(arg_params=arg, aux_params=aux,
                       allow_missing=True, allow_extra=True)

        '''Run inference on the image'''
        mod.forward(Batch([mx.nd.array(image_data)]))
        predictions = mod.get_outputs()[0].asnumpy()
        predictions = predictions[0].tolist()

        '''Apply emotion labels'''
        zipb_object = zip(self.__emotion_labels, predictions)
        prediction_dictionary = dict(zipb_object)

        return prediction_dictionary

预期的行为是异步运行每个图像的推理，但也等待整个批次的过程完成。

Answer 1

您不应该做的一件事是为每张图片加载模型。应该加载一次模型，然后对所有 600 张图像进行推理。

例如，您可以像这样重构代码：

def load_model(self):
        '''Load the model'''
        model_metadata = onnx_mxnet.get_model_metadata(self.__model)
        data_names = [inputs[0]
                      for inputs in model_metadata.get('input_tensor_data')]
        Batch = namedtuple('Batch', 'data')

        ctx = mx.eia()  # Set the context to elastic inference

        '''Load the model'''
        sym, arg, aux = onnx_mxnet.import_model(self.__model)
        mod = mx.mod.Module(symbol=sym, data_names=data_names,
                            context=ctx, label_names=None)
        mod.bind(data_shapes=[(data_names[0], image_data.shape)],
                 label_shapes=None, for_training=False)

        mod.set_params(arg_params=arg, aux_params=aux,
                       allow_missing=True, allow_extra=True)

        return mod


def infer(self, mod, target_image_path):
        target_image_path = self.__output_directory + '/' + target_image_path

        image_data = self.__get_image_data(target_image_path)  # Get pixel data

        '''Run inference on the image'''
        mod.forward(Batch([mx.nd.array(image_data)]))
        predictions = mod.get_outputs()[0].asnumpy()
        predictions = predictions[0].tolist()

        '''Apply emotion labels'''
        zipb_object = zip(self.__emotion_labels, predictions)
        prediction_dictionary = dict(zipb_object)

        return prediction_dictionary

MXNet 在异步引擎上运行，您无需等待图像完成处理即可将新图像加入队列。

MXNet 中的某些调用是异步的，例如当您调用mod.forward() 时，此调用会立即返回，并且不会等待计算结果。其他调用是同步的，例如mod.get_outputs()[0].asnumpy() 这会将数据复制到 CPU，因此它必须是同步的。在每次迭代之间进行同步调用会稍微减慢处理速度。

假设您可以访问 image_paths 列表，您可以像这样处理它们以最大限度地减少等待时间并仅在末尾有一个同步点：

    results = []
    for target_image_path in image_paths:
        image_data = self.__get_image_data(target_image_path)  # Get pixel data

        '''Run inference on the image'''
        mod.forward(Batch([mx.nd.array(image_data)]))
        results.append(mod.get_outputs()[0])
    predictions = [result.asnumpy()[0].tolist() for result in results]

您可以在此处阅读有关使用 MXNet 进行异步编程的更多信息：http://d2l.ai/chapter_computational-performance/async-computation.html

如果您知道要处理 N 个图像，则更好的是，您可以将它们分批处理，例如 16 个，以提高处理的并行度。但是这样做会增加内存消耗。由于您似乎使用的是弹性推理上下文，因此您的整体内存将受到限制，我建议您坚持使用较小的批量大小，以免内存不足。