Visión general
Esta página describe el diseño y los pasos necesarios para convertir operaciones compuestas en TensorFlow en operaciones fusionadas en TensorFlow Lite. Esta infraestructura es de propósito general y admite la conversión de cualquier operación compuesta en TensorFlow a una operación fusionada correspondiente en TensorFlow Lite.
Un ejemplo de uso de esta infraestructura es la fusión de la operación RNN de TensorFlow con TensorFlow Lite, como se detalla aquí .
¿Qué son las operaciones fusionadas?
Las operaciones de TensorFlow pueden ser operaciones primitivas, por ejemplo , tf.add , o pueden estar compuestas de otras operaciones primitivas, por ejemplo, tf.einsum . Una operación primitiva se muestra como un solo nodo en el gráfico de TensorFlow, mientras que una operación compuesta es una colección de nodos en el gráfico de TensorFlow. Ejecutar una operación compuesta es equivalente a ejecutar cada una de sus operaciones primitivas constituyentes.
Una operación fusionada corresponde a una sola operación que subsume todo el cálculo realizado por cada operación primitiva dentro de la operación compuesta correspondiente.
Beneficios de las operaciones fusionadas
Las operaciones fusionadas existen para maximizar el rendimiento de sus implementaciones de kernel subyacentes, al optimizar el cálculo general y reducir el consumo de memoria. Esto es muy valioso, especialmente para cargas de trabajo de inferencia de baja latencia y plataformas móviles con recursos limitados.
Las operaciones fusionadas también proporcionan una interfaz de nivel superior para definir transformaciones complejas como la cuantificación, que de otro modo sería inviable o muy difícil de realizar a un nivel más granular.
TensorFlow Lite tiene muchas instancias de operaciones fusionadas por las razones expuestas anteriormente. Estas operaciones fusionadas generalmente corresponden a operaciones compuestas en el programa fuente de TensorFlow. Los ejemplos de operaciones compuestas en TensorFlow que se implementan como una sola operación fusionada en TensorFlow Lite incluyen varias operaciones RNN como LSTM de secuencia unidireccional y bidireccional, convolución (conv2d, bias add, relu), totalmente conectado (matmul, bias add, relu) y más . En TensorFlow Lite, la cuantificación de LSTM actualmente solo se implementa en las operaciones de LSTM fusionadas.
Desafíos con operaciones fusionadas
Convertir operaciones compuestas de TensorFlow en operaciones fusionadas en TensorFlow Lite es un problema difícil. Esto es porque:
Las operaciones compuestas se representan en el gráfico de TensorFlow como un conjunto de operaciones primitivas sin un límite bien definido. Puede ser muy difícil identificar (por ejemplo, a través de la coincidencia de patrones) el subgráfico correspondiente a una operación compuesta de este tipo.
Puede haber más de una implementación de TensorFlow dirigida a una operación fusionada de TensorFlow Lite. Por ejemplo, hay muchas implementaciones de LSTM en TensorFlow (Keras, Babelfish/lingvo, etc.) y cada una de ellas se compone de diferentes operaciones primitivas, pero todas podrían convertirse a la misma operación LSTM fusionada en TensorFlow Lite.
Como tal, la conversión de operaciones fusionadas ha demostrado ser todo un desafío.
Conversión de una operación compuesta a una operación personalizada TFLite (recomendado)
Envuelva la operación compuesta en una tf.function
En muchos casos, una parte del modelo se puede asignar a una sola operación en TFLite. Esto puede ayudar con el rendimiento al escribir una implementación optimizada para operaciones específicas. Para poder crear una operación fusionada en TFLite, identifique la parte del gráfico que representa una operación fusionada y envuélvala en una tf.function con el atributo "experimental_implements" a una tf.function , que tiene el valor del atributo tfl_fusable_op con el valor true . Si la operación personalizada toma atributos, páselos como parte de los mismos "experimental_implements".
Ejemplo,
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))
Tenga en cuenta que no necesita configurar allow_custom_ops en el convertidor, ya que el atributo tfl_fusable_op ya lo implica.
Implemente operaciones personalizadas y regístrese con TFLite Interpreter
Implemente su operación fusionada como una operación personalizada de TFLite; consulte las instrucciones .
Tenga en cuenta que el nombre con el que registrar la operación debe ser similar al nombre especificado en el atributo de name en la firma del implemento.
Un ejemplo para el op en el ejemplo es
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, ®);
Conversión de operación compuesta a fusionada (avanzado)
La arquitectura general para convertir operaciones compuestas de TensorFlow en operaciones fusionadas de TensorFlow Lite se muestra a continuación:
Envuelva la operación compuesta en una tf.function
En el código fuente del modelo TensorFlow, identifica y abstrae la operación compuesta en una tf.function con la anotación de la función experimental_implements . Vea un ejemplo de búsqueda incrustada . La función define la interfaz y sus argumentos deben usarse para implementar la lógica de conversión.
Escribir código de conversión
El código de conversión se escribe según la interfaz de la función con la anotación de implements . Vea un ejemplo de fusión para la búsqueda incrustada . Conceptualmente, el código de conversión reemplaza la implementación compuesta de esta interfaz por la fusionada.
En el paso de preparación de funciones compuestas, inserte su código de conversión .
En usos más avanzados, es posible implementar transformaciones complejas de los operandos de la operación compuesta para derivar los operandos de la operación fusionada. Ver Keras LSTM . código de conversión como ejemplo.
Convertir a TensorFlow Lite
Usa la API TFLiteConverter.from_saved_model para convertir a TensorFlow Lite.
Bajo el capó
Ahora describimos detalles de alto nivel del diseño general en la conversión a operaciones fusionadas en TensorFlow Lite.
Composición de operaciones en TensorFlow
El uso de tf.function con el atributo de función experimental_implements permite a los usuarios crear explícitamente nuevas operaciones mediante operaciones primitivas de TensorFlow y especificar la interfaz que implementa la operación compuesta resultante. Esto es muy útil ya que proporciona:
- Un límite bien definido para la operación compuesta en el gráfico de TensorFlow subyacente.
- Especifique explícitamente la interfaz que implementa esta operación. Los argumentos de la
tf.functioncorresponden a los argumentos de esta interfaz.
Como ejemplo, consideremos una operación compuesta definida para implementar la búsqueda incrustada. Esto se asigna a una operación fusionada en TensorFlow Lite.
@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
Al hacer que los modelos usen operaciones compuestas a través de tf.function como se ilustra arriba, es posible construir una infraestructura general para identificar y convertir dichas operaciones en operaciones fusionadas de TensorFlow Lite.
Ampliación del convertidor TensorFlow Lite
El convertidor TensorFlow Lite que se lanzó a principios de este año solo admitía la importación de modelos de TensorFlow como un gráfico con todas las variables reemplazadas por sus valores constantes correspondientes. Esto no funciona para la fusión de operaciones ya que dichos gráficos tienen todas las funciones en línea para que las variables se puedan convertir en constantes.
Para aprovechar la tf.function con la característica experimental_implements durante el proceso de conversión, las funciones deben conservarse hasta más adelante en el proceso de conversión.
Como tal, implementamos un nuevo flujo de trabajo de importación y conversión de modelos de TensorFlow en el convertidor para admitir el caso de uso de fusión de operaciones compuestas. En concreto, las novedades añadidas son:
- Importación de modelos guardados de TensorFlow en MLIR
- fusionar operaciones compuestas
- análisis de mutabilidad variable
- congelar todas las variables de solo lectura
Esto nos permite realizar la fusión de operaciones utilizando las funciones que representan las operaciones compuestas antes de la incorporación de funciones y la congelación de variables.
Implementación de la operación fusión
Veamos la operación pase de fusión con más detalle. Este pase hace lo siguiente:
- Recorra todas las funciones en el módulo MLIR.
- Si una función tiene el atributo tf._implements, según el valor del atributo, llama a la utilidad de fusión de operación adecuada.
- La utilidad de fusión de operaciones opera en los operandos y atributos de la función (que sirven como interfaz para la conversión) y reemplaza el cuerpo de la función con un cuerpo de función equivalente que contiene la operación fusionada.
- En muchos casos, el cuerpo reemplazado contendrá operaciones distintas a la operación fusionada. Estos corresponden a algunas transformaciones estáticas en los operandos de la función para obtener los operandos de la operación fusionada. Dado que todos estos cálculos se pueden plegar constantemente, no estarían presentes en el búfer plano exportado donde solo existiría la operación fusionada.
Aquí hay un fragmento de código del pase que muestra el flujo de trabajo principal:
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 */
}
Aquí hay un fragmento de código que muestra la asignación de esta operación compuesta a una operación fusionada en TensorFlow Lite aprovechando la función como una interfaz de conversión.
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());
}