Как я могу сделать простой matmul на краю тпу? - PullRequest
Новая
       29

Как я могу сделать простой matmul на краю тпу?

0 голосов
/ 03 июня 2019

Я не могу понять, как вызвать мою модель .tflite, которая делает matmul на ускорителе кораллов, используя API Python.

Модель .tflite генерируется из некоторого примера кода здесь . Он хорошо работает с использованием класса tf.lite.Interpreter (), но я не знаю, как преобразовать его для работы с классом edgetpu. Я попробовал edgetpu.basic.basic_engine.BasicEngine (), изменив тип данных моделей с numpy.float32 на numpy.uint8, но это не помогло. Я новичок в TensorFlow и просто хочу использовать мой ТПУ для матмуля. 

import numpy
import tensorflow as tf
import edgetpu
from edgetpu.basic.basic_engine import BasicEngine

def export_tflite_from_session(session, input_nodes, output_nodes, tflite_filename):
    print("Converting to tflite...")
    converter = tf.lite.TFLiteConverter.from_session(session, input_nodes, output_nodes)
    tflite_model = converter.convert()
    with open(tflite_filename, "wb") as f:
        f.write(tflite_model)
    print("Converted %s." % tflite_filename)

#This does matmul just fine but does not use the TPU
def test_tflite_model(tflite_filename, examples):
    print("Loading TFLite interpreter for %s..." % tflite_filename)
    interpreter = tf.lite.Interpreter(model_path=tflite_filename)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    print("input details: %s" % input_details)
    print("output details: %s" % output_details)

    for i, input_tensor in enumerate(input_details):
        interpreter.set_tensor(input_tensor['index'], examples[i])
    interpreter.invoke()
    model_output = []
    for i, output_tensor in enumerate(output_details):
        model_output.append(interpreter.get_tensor(output_tensor['index']))
    return model_output

#this should use the TPU, but I don't know how to run the model or if it needs
#further processing. One matrix can be constant for my use case
def test_tpu(tflite_filename,examples):
    print("Loading TFLite interpreter for %s..." % tflite_filename)
    #TODO edgetpu.basic
    interpreter = BasicEngine(tflite_filename)
    interpreter.allocate_tensors()#does not work...

def main():
    tflite_filename = "model.tflite"
    shape_a = (2, 2)
    shape_b = (2, 2)

    a = tf.placeholder(dtype=tf.float32, shape=shape_a, name="A")
    b = tf.placeholder(dtype=tf.float32, shape=shape_b, name="B")
    c = tf.matmul(a, b, name="output")

    numpy.random.seed(1234)
    a_ = numpy.random.rand(*shape_a).astype(numpy.float32)
    b_ = numpy.random.rand(*shape_b).astype(numpy.float32)
    with tf.Session() as session:
        session_output = session.run(c, feed_dict={a: a_, b: b_})
        export_tflite_from_session(session, [a, b], [c], tflite_filename)

    tflite_output = test_tflite_model(tflite_filename, [a_, b_])
    tflite_output = tflite_output[0]

    #test the TPU
    tflite_output = test_tpu(tflite_filename, [a_, b_])

    print("Input example:")
    print(a_)
    print(a_.shape)
    print(b_)
    print(b_.shape)
    print("Session output:")
    print(session_output)
    print(session_output.shape)
    print("TFLite output:")
    print(tflite_output)
    print(tflite_output.shape)
    print(numpy.allclose(session_output, tflite_output))

if __name__ == '__main__':
    main()

1 Ответ

0 голосов
/ 05 июня 2019

Вы конвертируете свою модель только один раз, и ваша модель не полностью скомпилирована для Edge TPU. Из документов :

В первой точке графа модели, где происходит неподдерживаемая операция, компилятор разбивает график на две части. Первая часть графика, которая содержит только поддерживаемые операции, скомпилирована в пользовательскую операцию, которая выполняется на Edge TPU, а все остальное выполняется на CPU

Существует несколько особых требований, которым должна соответствовать модель:

  • обучение с учетом квантования

  • постоянные размеры тензора и параметры модели во время компиляции

  • тензоры 3-мерные или меньше.

  • используют только операции, поддерживаемые Edge TPU.

Существует онлайн-компилятор , а также CLI-версия , которая полезна для перевода моделей .tflite в Edge TPU-совместимые модели .tflite.

Ваш код также неполный. Вы передали свою модель в класс здесь:

interpreter = BasicEngine(tflite_filename)

но вам не хватает шага, чтобы фактически выполнить логический вывод для тензора: 

output = RunInference(interpreter)
...