Обзор
На этой странице описывается структура и шаги, необходимые для преобразования составных операций TensorFlow в объединённые операции LiteRT. Эта инфраструктура является универсальной и поддерживает преобразование любой составной операции TensorFlow в соответствующую объединённую операцию LiteRT.
Примером использования этой инфраструктуры является слияние операций TensorFlow RNN с LiteRT, как подробно описано здесь .
Что такое объединенные операции?
Операции TensorFlow могут быть примитивными операциями, например, tf.add , или составлены из других примитивных операций, например, tf.einsum . Примитивная операция отображается как отдельный узел в графе TensorFlow, а составная операция — как набор узлов в графе TensorFlow. Выполнение составной операции эквивалентно выполнению каждой из входящих в неё примитивных операций.
Объединенная операция соответствует одной операции, которая включает все вычисления, выполняемые каждой примитивной операцией, в соответствующую составную операцию.
Преимущества объединенных операций
Объединенные операции предназначены для максимального повышения производительности базовых реализаций ядра за счет оптимизации общих вычислений и сокращения потребления памяти. Это очень ценно, особенно для задач логического вывода с низкой задержкой и мобильных платформ с ограниченными ресурсами.
Объединенные операции также предоставляют интерфейс более высокого уровня для определения сложных преобразований, таких как квантование, которые в противном случае было бы невозможно или очень сложно реализовать на более детальном уровне.
В LiteRT существует множество примеров объединённых операций по причинам, изложенным выше. Эти объединённые операции обычно соответствуют составным операциям в исходной программе TensorFlow. Примеры составных операций в TensorFlow, реализованных в LiteRT как одна объединённая операция, включают различные операции RNN, такие как однонаправленные и двунаправленные последовательности LSTM, свёртки (conv2d, bias add, relu), полносвязные (matmul, bias add, relu) и другие. В LiteRT квантование LSTM в настоящее время реализовано только в объединённых операциях LSTM.
Проблемы с объединенными операциями
Преобразование составных операций из TensorFlow в объединённые операции в LiteRT — сложная задача. Это связано со следующими причинами:
Составные операции представлены в графе TensorFlow как набор примитивных операций без чётко определённых границ. Определить (например, с помощью сопоставления с образцом) подграф, соответствующий такой составной операции, может быть очень сложно.
Может существовать несколько реализаций TensorFlow, ориентированных на объединённую операцию LiteRT. Например, в TensorFlow существует множество реализаций LSTM (Keras, Babelfish/lingvo и т. д.), и каждая из них состоит из различных примитивных операций, но все они могут быть преобразованы в одну и ту же объединённую операцию LSTM в LiteRT.
Таким образом, преобразование объединенных операций оказалось весьма сложной задачей.
Преобразование составной операции в пользовательскую операцию TFLite (рекомендуется)
Оберните составную операцию в tf.function
Во многих случаях часть модели можно сопоставить с одной операцией в TFLite. Это может повысить производительность при написании оптимизированной реализации для конкретных операций. Чтобы создать объединённую операцию в TFLite, определите часть графа, представляющую объединённую операцию, и оберните её в функцию tf.function с атрибутом "experimental_implements" в функцию tf.function , имеющую значение атрибута tfl_fusable_op со значением true . Если пользовательская операция принимает атрибуты, передайте их как часть того же "experimental_implements".
Пример,
def get_implements_signature():
implements_signature = [
# 'name' will be used as a name for the operation.
'name: "my_custom_fused_op"',
# attr "tfl_fusable_op" is required to be set with true value.
'attr {key: "tfl_fusable_op" value { b: true } }',
# Example attribute "example_option" that the op accepts.
'attr {key: "example_option" value { i: %d } }' % 10
]
return ' '.join(implements_signature)
@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
# An empty function that represents pre/post processing example that
# is not represented as part of the Tensorflow graph.
output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
return output_1, output_2
class TestModel(tf.Module):
def __init__(self):
super(TestModel, self).__init__()
self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
@tf.function(input_signature=[
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
])
def simple_eval(self, input_a, input_b):
return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))
Обратите внимание, что вам не нужно устанавливать allow_custom_ops для конвертера, так как атрибут tfl_fusable_op уже подразумевает это.
Реализуйте пользовательскую операцию и зарегистрируйте ее с помощью TFLite Interpreter
Реализуйте объединенную операцию как пользовательскую операцию TFLite — см. инструкции .
Обратите внимание, что имя, под которым регистрируется операция, должно быть похоже на имя, указанное в атрибуте name в сигнатуре implements.
Пример для операции в примере:
TfLiteRegistration reg = {};
// This name must match the name specified in the implements signature.
static constexpr char kOpName[] = "my_custom_fused_op";
reg.custom_name = kOpName;
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, ®);
Переход от композитной к слитной операции (расширенная)
Общая архитектура преобразования составных операций TensorFlow в объединенные операции LiteRT представлена ниже:
Оберните составную операцию в tf.function
В исходном коде модели TensorFlow определите и абстрагируйте составную операцию в функцию tf.function с помощью аннотации experimental_implements . См. пример встраивания lookup . Функция определяет интерфейс, а её аргументы должны использоваться для реализации логики преобразования.
Написать код преобразования
Код преобразования написан на основе интерфейса функции с аннотацией implements . См. пример слияния для поиска встраивания . Концептуально, код преобразования заменяет составную реализацию этого интерфейса на слитую.
В проходе prepare-composite-functions подключите код преобразования .
В более сложных случаях можно реализовать сложные преобразования операндов составной операции для получения операндов объединённой операции. См. пример кода преобразования Keras LSTM .
Конвертировать в LiteRT
Используйте API TFLiteConverter.from_saved_model для конвертации в LiteRT.
Под капотом
Теперь мы опишем общие детали общей конструкции при переходе к объединенным операциям в LiteRT.
Составление операций в TensorFlow
Использование tf.function с атрибутом experimental_implements позволяет пользователям явно создавать новые операции с использованием примитивных операций TensorFlow и указывать интерфейс, который реализует полученная составная операция. Это очень полезно, поскольку обеспечивает:
- Четко определенная граница для составной операции в базовом графе TensorFlow.
- Явно укажите интерфейс, реализуемый этой операцией. Аргументы функции tf.function соответствуют аргументам этого интерфейса.
В качестве примера рассмотрим составную операцию, определённую для реализации поиска по встраиваемым данным. Она соответствует объединённой операции в LiteRT.
@tf.function(
experimental_implements="embedding_lookup")
def EmbFprop(embs, ids_vec):
"""Embedding forward prop.
Effectively, it computes:
num = size of ids_vec
rets = zeros([num, embedding dim])
for i in range(num):
rets[i, :] = embs[ids_vec[i], :]
return rets
Args:
embs: The embedding matrix.
ids_vec: A vector of int32 embedding ids.
Returns:
The result of embedding lookups. A matrix of shape
[num ids in ids_vec, embedding dims].
"""
num = tf.shape(ids_vec)[0]
rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))
def EmbFpropLoop(i, embs, ids_vec, rets):
# row_id = ids_vec[i]
row_id = tf.gather(ids_vec, i)
# row = embs[row_id]
row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
# rets[i] = row
rets = inplace_ops.alias_inplace_update(rets, [i], row)
return embs, ids_vec, rets
_, _, rets = functional_ops.For(
start=0,
limit=num,
delta=1,
inputs=[embs, ids_vec, rets],
body=EmbFpropLoop,
rewrite_with_while=compiled)
if len(weight_shape) > 2:
rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
return rets
Заставляя модели использовать составные операции через tf.function , как показано выше, становится возможным построить общую инфраструктуру для идентификации и преобразования таких операций в объединенные операции LiteRT.
Расширение конвертера LiteRT
Конвертер LiteRT, выпущенный ранее в этом году, поддерживал импорт моделей TensorFlow только в виде графа, в котором все переменные заменялись соответствующими константами. Это не работает для операций слияния, поскольку в таких графах все функции встроены, что позволяет преобразовать переменные в константы.
Чтобы использовать tf.function с функцией experimental_implements в процессе преобразования, функции необходимо сохранить до более поздней стадии процесса преобразования.
В связи с этим мы реализовали новый рабочий процесс импорта и преобразования моделей TensorFlow в конвертере для поддержки варианта использования композитных операций. В частности, добавлены следующие новые функции:
- Импорт сохраненных моделей TensorFlow в MLIR
- композитные операции слияния
- анализ изменчивости переменных
- заморозить все переменные, доступные только для чтения
Это позволяет нам выполнять слияние операций с использованием функций, представляющих составные операции, до встраивания функций и заморозки переменных.
Реализация операции «Слияние»
Давайте рассмотрим проход операции «слияние» более подробно. Этот проход выполняет следующие действия:
- Выполните цикл по всем функциям в модуле MLIR.
- Если функция имеет атрибут tf._implements, на основе значения атрибута вызывает соответствующую утилиту слияния операций.
- Утилита слияния операций работает с операндами и атрибутами функции (которые служат интерфейсом для преобразования) и заменяет тело функции эквивалентным телом функции, содержащим объединенную операцию.
- Во многих случаях заменяемое тело будет содержать операции, отличные от операции объединённой. Они соответствуют некоторым статическим преобразованиям операндов функции для получения операндов объединённой операции. Поскольку все эти вычисления можно свернуть с помощью констант, они не будут присутствовать в экспортированном плоском буфере, где будет существовать только операция объединённой.
Вот фрагмент кода из прохода, показывающий основной рабочий процесс:
void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
StringAttr attr) {
if (attr.getValue() == "embedding_lookup") {
func.eraseBody();
func.addEntryBlock();
// Convert the composite embedding_lookup function body to a
// TFLite fused embedding_lookup op.
ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
if (failed(convert_embedded_lookup.VerifySignature())) {
return signalPassFailure();
}
convert_embedded_lookup.RewriteFunc();
} else if (attr.getValue() == mlir::TFL::kKerasLstm) {
func.eraseBody();
func.addEntryBlock();
OpBuilder builder(func.getBody());
if (failed(ConvertKerasLSTMLayer(func, &builder))) {
return signalPassFailure();
}
} else if (.....) /* Other fusions can plug in here */
}
Ниже представлен фрагмент кода, демонстрирующий сопоставление этой составной операции с объединенной операцией в LiteRT, использующей функцию как интерфейс преобразования.
void RewriteFunc() {
Value lookup = func_.getArgument(1);
Value value = func_.getArgument(0);
auto output_type = func_.getType().getResult(0);
OpBuilder builder(func_.getBody());
auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
func_.getLoc(), output_type, lookup, value);
builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
}