Descripción general
En esta página, se describen el diseño y los pasos necesarios para convertir operaciones compuestas en TensorFlow en operaciones fusionadas en TensorFlow Lite. Esta infraestructura es de uso general y admite la conversión de cualquier operación compuesta en TensorFlow en una operación fusionada correspondiente en TensorFlow Lite.
Un ejemplo de uso de esta infraestructura es la fusión de operaciones de RNN de TensorFlow con TensorFlow Lite, como se detalla aquí.
Qué son las operaciones fusionadas
Las operaciones de TensorFlow pueden ser operaciones primitivas, p.ej., tf.add, o pueden estar compuestas a partir de otras operaciones primitivas, p.ej., tf.einsum. Una operación básica se muestra como un nodo único en el grafo de TensorFlow, mientras que una operación compuesta es una colección de nodos en el grafo de TensorFlow. Ejecutar una operación compuesta equivale a ejecutar cada una de las operaciones primitivas que constituyen.
Una operación fusionada corresponde a una operación única que resume 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 las implementaciones de kernel subyacentes, ya que optimiza el procesamiento general y reduce el espacio en memoria. Esto es muy valioso, en especial 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 cuantización, que, de otro modo, serían inviables o muy difíciles de hacer con un nivel más detallado.
TensorFlow Lite tiene muchas instancias de operaciones fusionadas por los motivos articulados anteriormente. Por lo general, estas operaciones fusionadas corresponden a operaciones compuestas en el programa de TensorFlow de origen. Algunos ejemplos de operaciones compuestas en TensorFlow que se implementan como una sola operación fusionada en TensorFlow Lite incluyen varias operaciones de RNN como LSTM de secuencia unidireccional y bidireccional, convolución (conv2d, sesgo agregado, relu), completamente conectadas (matmul, sesgo add, Relu) y mucho más. En TensorFlow Lite, actualmente, la cuantización de LSTM solo se implementa en las operaciones de LSTM fusionadas.
Desafíos de las operaciones fusionadas
Convertir operaciones compuestas de TensorFlow en operaciones fusionadas en TensorFlow Lite es un problema difícil. Esto se debe a los siguientes motivos:
Las operaciones compuestas se representan en el grafo de TensorFlow como un conjunto de operaciones primitivas sin un límite bien definido. Puede ser muy difícil identificar (p.ej., a través de la coincidencia de patrones) el subgrafo correspondiente a esa operación compuesta.
Es posible que haya más de una implementación de TensorFlow orientada 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 está compuesta por diferentes operaciones primitivas, pero todas se pueden convertir en la misma operación de LSTM fusionada en TensorFlow Lite.
Por ello, la conversión de operaciones fusionadas ha demostrado ser bastante desafiante.
Pasa de una op compuesta a una operación personalizada de TFLite (recomendado)
Une la operación compuesta en un tf.function
En muchos casos, alguna parte del modelo se puede asignar a una sola operación en TFLite. Esto puede ayudar con el rendimiento cuando escribes una implementación optimizada para operaciones específicas. Para poder crear una operación fusionada en TFLite, identifica la parte del gráfico que representa una operación fusionada y únela en una tf.function con el atributo "experimental_implements" en una tf.function, que tenga el valor de atributo tfl_fusable_op con el valor true. Si la operación personalizada toma atributos, pásalos como parte del mismo "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))
Ten en cuenta que no necesitas configurar allow_custom_ops en el convertidor, ya que el atributo tfl_fusable_op ya lo implica.
Implementa operaciones personalizadas y regístrate con el intérprete de TFLite
Implementa tu operación combinada como una operación personalizada de TFLite. Consulta las instructions.
Ten en cuenta que el nombre con el que se registrará la operación debe ser similar al nombre especificado en el atributo name en la firma de la implementación.
Un ejemplo de la op del 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, ®);
Pasa de una operación compuesta a una fusionada (avanzado)
A continuación, se muestra la arquitectura general para convertir las operaciones compuestas de TensorFlow en operaciones fusionadas de TensorFlow Lite:
Une la operación compuesta en un tf.function
En el código fuente del modelo de TensorFlow, identifica y abstrae la operación compuesta en una tf.function con la anotación de función experimental_implements. Consulta un ejemplo de búsqueda de incorporación. 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 implements. Consulta un ejemplo de fusión para la búsqueda de incorporaciones. De forma conceptual, el código de conversión reemplaza la implementación compuesta de esta interfaz por la fusionada.
En el pase de prepare-composite-functions, agrega tu 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. Consulta el código de conversión KerasTM como ejemplo.
Cámbiate a TensorFlow Lite
Usa la API TFLiteConverter.from_saved_model para convertir a TensorFlow Lite.
Detrás de escena
Ahora describimos los detalles de alto nivel del diseño general para convertir en operaciones fusionadas en TensorFlow Lite.
Operaciones de composición en TensorFlow
El uso de tf.function con el atributo de función experimental_implements permite a los usuarios redactar de forma explícita operaciones nuevas con operaciones básicas de TensorFlow y especificar la interfaz que implementa la operación compuesta resultante. Esto es muy útil, ya que proporciona lo siguiente:
- Un límite bien definido para la operación compuesta en el grafo de TensorFlow subyacente.
- Especifica de forma explícita la interfaz que implementa esta operación. Los argumentos de tf.function corresponden a los argumentos de esta interfaz.
A modo de ejemplo, consideremos una operación compuesta definida para implementar una búsqueda de incorporación. Esto se asigna a una operación combinada 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
Cuando haces que los modelos usen operaciones compuestas a través de tf.function como se ilustra más arriba, se puede crear una infraestructura general para identificar y convertir esas operaciones en operaciones fusionadas de TensorFlow Lite.
Extiende el conversor de TensorFlow Lite
El conversor de 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, en el que todas las variables se reemplazaron por sus valores constantes correspondientes. Esto no funciona para la fusión de operaciones, ya que esos gráficos tienen todas las funciones intercaladas de modo que las variables se puedan convertir en constantes.
Para aprovechar tf.function con la función experimental_implements durante el proceso de conversión, las funciones deben conservarse hasta una etapa posterior en el proceso de conversión.
Por lo tanto, implementamos un nuevo flujo de trabajo de importación y conversión de modelos de TensorFlow en el convertidor para admitir el caso práctico de fusión de operaciones compuestas. En concreto, las nuevas funciones que se agregan son las siguientes:
- Importación de modelos guardados de TensorFlow a MLIR
- fusionar operaciones compuestas
- análisis de mutabilidad de variables
- inmovilizar todas las variables de solo lectura
Esto nos permite realizar la fusión de operaciones con las funciones que representan las operaciones compuestas antes de la intercalación de funciones y la congelación de variables.
Implementa la fusión de operaciones
Veamos con más detalle el paso de fusión de la operación. Este pase hace lo siguiente:
- Recorre todas las funciones en el módulo de MLIR.
- Si una función tiene el atributo tf._implements, según el valor del atributo, llama a la utilidad de fusión de operaciones adecuada.
- La utilidad de fusión de operaciones opera en los operandos y atributos de la función (que funcionan como interfaz para la conversión) y reemplaza el cuerpo de la función por uno equivalente que contiene la operación fusionada.
- En muchos casos, el cuerpo reemplazado contendrá operaciones distintas de la operación combinada. Estas 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 de forma constante, no estarían presentes en el búfer plano exportado, donde solo existiría la operación fusionada.
Este es 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 */
}
Este es un fragmento de código en el que se muestra la asignación de esta operación compuesta a una operación fusionada en TensorFlow Lite que aprovecha 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());
}