TFX для мобильных устройств

Введение

В этом руководстве показано, как Tensorflow Extended (TFX) может создавать и оценивать модели машинного обучения, которые будут развернуты на устройстве. TFX теперь обеспечивает встроенную поддержку TFLite , что позволяет выполнять высокоэффективный вывод на мобильных устройствах.

В этом руководстве рассказывается об изменениях, которые можно внести в любой конвейер для создания и оценки моделей TFLite. Здесь мы приводим полный пример, демонстрирующий, как TFX может обучать и оценивать модели TFLite, обученные на наборе данных MNIST . Кроме того, мы покажем, как один и тот же конвейер можно использовать для одновременного экспорта как стандартной модели SavedModel на основе Keras, так и модели TFLite, что позволяет пользователям сравнивать качество двух моделей.

Мы предполагаем, что вы знакомы с TFX, нашими компонентами и нашими конвейерами. Если нет, то, пожалуйста, посмотрите этот урок .

Шаги

Для создания и оценки модели TFLite в TFX требуется всего два шага. Первым шагом является вызов переписчика TFLite в контексте TFX Trainer для преобразования обученной модели TensorFlow в модель TFLite. Второй шаг — настройка Evaluator для оценки моделей TFLite. Теперь мы обсудим каждый по очереди.

Вызов переписчика TFLite в трейнере.

Тренер TFX ожидает, что определяемый пользователем run_fn будет указан в файле модуля. Этот run_fn определяет модель, подлежащую обучению, обучает ее указанному количеству итераций и экспортирует обученную модель.

В оставшейся части этого раздела мы предоставляем фрагменты кода, которые показывают изменения, необходимые для вызова переписчика TFLite и экспорта модели TFLite. Весь этот код находится в run_fn модуля MNIST TFLite .

Как показано в коде ниже, мы должны сначала создать подпись, которая принимает Tensor для каждой функции в качестве входных данных. Обратите внимание, что это отличие от большинства существующих моделей в TFX, которые в качестве входных данных принимают сериализованные прототипы tf.Example .

 signatures = {
      'serving_default':
          _get_serve_tf_examples_fn(
              model, tf_transform_output).get_concrete_function(
                  tf.TensorSpec(
                      shape=[None, 784],
                      dtype=tf.float32,
                      name='image_floats'))
  }

Затем модель Keras сохраняется как SavedModel, как обычно.

  temp_saving_model_dir = os.path.join(fn_args.serving_model_dir, 'temp')
  model.save(temp_saving_model_dir, save_format='tf', signatures=signatures)

Наконец, мы создаем экземпляр переписчика TFLite ( tfrw ) и вызываем его в SavedModel, чтобы получить модель TFLite. Мы сохраняем эту модель TFLite в serving_model_dir , предоставленном вызывающей стороной run_fn . Таким образом, модель TFLite сохраняется в том месте, где все последующие компоненты TFX ожидают найти модель.

  tfrw = rewriter_factory.create_rewriter(
      rewriter_factory.TFLITE_REWRITER, name='tflite_rewriter')
  converters.rewrite_saved_model(temp_saving_model_dir,
                                 fn_args.serving_model_dir,
                                 tfrw,
                                 rewriter.ModelType.TFLITE_MODEL)

Оценка модели TFLite.

TFX Evaluator предоставляет возможность анализировать обученные модели, чтобы понять их качество по широкому спектру показателей. Помимо анализа сохраненных моделей, TFX Evaluator теперь также может анализировать модели TFLite.

В следующем фрагменте кода (воспроизведенном из конвейера MNIST ) показано, как настроить оценщик, который анализирует модель TFLite.

  # Informs the evaluator that the model is a TFLite model.
  eval_config_lite.model_specs[0].model_type = 'tf_lite'

  ...

  # Uses TFMA to compute the evaluation statistics over features of a TFLite
  # model.
  model_analyzer_lite = Evaluator(
      examples=example_gen.outputs['examples'],
      model=trainer_lite.outputs['model'],
      eval_config=eval_config_lite,
  ).with_id('mnist_lite')

Как показано выше, единственное изменение, которое нам нужно сделать, — это установить для поля model_type значение tf_lite . Для анализа модели TFLite не требуется никаких других изменений конфигурации. Независимо от того, анализируется ли модель TFLite или SavedModel, выходные данные Evaluator будут иметь одинаковую структуру.

Однако обратите внимание, что оценщик предполагает, что модель TFLite сохранена в файле с именем tflite в файле train_lite.outputs['model'].