Преобразование из замороженного тензорного графа в tflite для вывода android - PullRequest
Новая
       35

Преобразование из замороженного тензорного графа в tflite для вывода android

0 голосов
/ 24 февраля 2020

Я пытаюсь преобразовать из pytorch в tflite для вывода android для приложения, над которым я работаю, которое использует данные камеры баскетбола в реальном времени для создания тепловой карты сделанных и пропущенных снимков. Это уже работает на iOS. Вот демонстрация .

Мне удалось преобразовать из pytorch (.pth) в onnx и из onnx в замороженный граф с тензорным потоком (.pb). Вывод на этот замороженный граф проверяется.

Однако, когда я пытаюсь преобразовать замороженный граф в tflite, я получаю следующую ошибку:

RuntimeError: Inputs and outputs not all float|uint8|int16 types.Node number 2 (ADD) failed to invoke.

Вводные данные от интерпретатора [interpreter.get_input_details (), interpreter.get_output_details ()] предполагают, что тип данных numpy .float32, и это то, где я запутался. Разве это не должно рассматриваться как плавание? Будем очень благодарны за любые предложения / помощь! 

[{'name': 'image', 'index': 21904, 'shape': array([  3, width, height], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]
[{'name': 'action', 'index': 7204, 'shape': array([], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]

Из документации тензорного потока: 

import numpy as np
import tensorflow as tf

# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="converted_model.tflite")
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Test model on random input data.
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)

interpreter.invoke()

# The function `get_tensor()` returns a copy of the tensor data.
# Use `tensor()` in order to get a pointer to the tensor.
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

Файл вывода для замороженного графа тензорного потока (1.14.0) для любого нуждающегося: 

import numpy as np
import tensorflow as tf
from PIL import Image

w = ...
h = ...

class CNN(object):

    def __init__(self, model_filepath):
        self.model_filepath = model_filepath
        self.load_graph(model_filepath = self.model_filepath)

    def load_graph(self, model_filepath):
        self.graph = tf.Graph()

        with tf.gfile.GFile(model_filepath, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())

        with self.graph.as_default():
            self.input = tf.placeholder(np.float32, shape = [3, h, w], name='image')
            tf.import_graph_def(graph_def, {'image': self.input})

        self.graph.finalize()
        self.sess = tf.Session(graph = self.graph)

    def test(self, data):
        output_tensor = self.graph.get_tensor_by_name('import/action:0')
        output = self.sess.run(output_tensor, feed_dict = {self.input: data})
        return output

def main():
    nn = CNN(model_filepath='out_1.14.pb')
    img = np.asarray(Image.open('example.jpg')).astype(np.float32)
    img = img.transpose(-1, 0, 1)
    ans = nn.test(data=img)
    print(ans)


if __name__ == '__main__':
    main()
...