Descripción general
En esta página, se describen el diseño y los pasos necesarios para convertir operaciones compuestas en TensorFlow a operaciones fusionadas en LiteRT. Esta infraestructura es de uso general y admite la conversión de cualquier operación compuesta en TensorFlow a una operación fusionada correspondiente en LiteRT.
Un ejemplo del uso de esta infraestructura es la fusión de operaciones de RNN de TensorFlow en LiteRT, como se detalla aquí.
Qué son las operaciones fusionadas
Las operaciones de TensorFlow pueden ser primitivas, p.ej., tf.add, o bien pueden componerse a partir de otras operaciones primitivas, p.ej., tf.einsum. Una operación primitiva aparece como un solo nodo 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 sus operaciones primitivas constituyentes.
Una operación fusionada corresponde a una sola operación que abarca todos los cálculos realizados 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, ya que optimizan el cálculo general y reducen el espacio en memoria. Esto es muy valioso, en especial para las cargas de trabajo de inferencia de baja latencia y las 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ían inviables o muy difíciles de realizar en un nivel más granular.
LiteRT tiene muchas instancias de operaciones fusionadas por los motivos que se mencionaron anteriormente. Por lo general, estas operaciones fusionadas corresponden a operaciones compuestas en el programa fuente de TensorFlow. Entre los ejemplos de operaciones compuestas en TensorFlow que se implementan como una sola operación fusionada en LiteRT, se incluyen varias operaciones de RNN, como LSTM de secuencia unidireccional y bidireccional, convolución (conv2d, bias add, relu), completamente conectada (matmul, bias add, relu) y muchas más. En LiteRT, la cuantificación de LSTM actualmente solo se implementa en las operaciones de LSTM fusionadas.
Desafíos con las operaciones fusionadas
Convertir operaciones compuestas de TensorFlow a operaciones fusionadas en LiteRT 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 una operación compuesta de este tipo.
Puede haber más de una implementación de TensorFlow que se dirija a una operación de LiteRT fusionada. 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 se pueden convertir a la misma operación de LSTM fusionada en LiteRT.
Por lo tanto, la conversión de operaciones fusionadas resultó ser bastante difícil.
Cómo convertir una operación compuesta en 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 se escribe una implementación optimizada para operaciones específicas. Para poder crear una operación fusionada en TFLite, identifica la parte del grafo que representa una operación fusionada y envuélvela en una tf.function con el atributo "experimental_implements" establecido en un tf.function, que tiene el valor del 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 es necesario que establezcas allow_custom_ops en el convertidor, ya que el atributo tfl_fusable_op ya lo implica.
Implementa una operación personalizada y regístrala con el intérprete de TFLite
Implementa tu operación fusionada como una operación personalizada de TFLite. Consulta las instrucciones.
Ten en cuenta que el nombre con el que se registra la operación debe ser similar al que se especifica en el atributo name de la firma de implementación.
Un ejemplo de la operación 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, ®);
Cómo convertir de una operación compuesta a una fusionada (Advanced)
A continuación, se muestra la arquitectura general para convertir operaciones compuestas de TensorFlow en operaciones fusionadas de LiteRT:
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 incorporaciones. La función define la interfaz y sus argumentos se deben usar para implementar la lógica de conversión.
Escribe el 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. Conceptualmente, el código de conversión reemplaza la implementación compuesta de esta interfaz por la fusionada.
En el paso prepare-composite-functions, conecta 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 de LSTM de Keras como ejemplo.
Cómo convertir a LiteRT
Usa la API de TFLiteConverter.from_saved_model para convertir el modelo a LiteRT.
Detrás de escena
Ahora describimos detalles de alto nivel del diseño general para convertir operaciones fusionadas en LiteRT.
Cómo componer operaciones en TensorFlow
El uso de tf.function con el atributo de función experimental_implements permite a los usuarios componer explícitamente nuevas operaciones con operaciones primitivas de TensorFlow y especificar la interfaz que implementa la operación compuesta resultante. Esto es muy útil, ya que proporciona lo siguiente:
- Es 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.
Como ejemplo, consideremos una operación compuesta definida para implementar la búsqueda de incorporaciones. Esto se asigna a una operación fusionada en 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
Si los modelos usan operaciones compuestas a través de tf.function, como se ilustró anteriormente, es posible crear una infraestructura general para identificar y convertir esas operaciones en operaciones fusionadas de LiteRT.
Cómo extender el convertidor de LiteRT
El conversor de LiteRT que se lanzó a principios de este año solo admitía la importación de modelos de TensorFlow como un grafo con todas las variables reemplazadas por sus valores constantes correspondientes. Esto no funciona para la fusión de operaciones, ya que esos grafos tienen todas las funciones intercaladas para 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 más adelante en el proceso de conversión.
Por lo tanto, implementamos un nuevo flujo de trabajo para importar y convertir modelos de TensorFlow en el convertidor para admitir el caso de uso de la fusión de operaciones compuestas. Específicamente, las nuevas funciones agregadas son las siguientes:
- Importación de modelos guardados de TensorFlow en MLIR
- fusionar operaciones compuestas
- análisis de mutabilidad de variables
- Inmoviliza 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 expansión de funciones y la congelación de variables.
Implementa la fusión de operaciones
Veamos el paso de fusión de operaciones con más detalle. Este pase hace lo siguiente:
- Itera todas las funciones del 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 sirven como interfaz para la conversión) y reemplaza el cuerpo de la función por un cuerpo de función equivalente que contiene la operación fusionada.
- En muchos casos, el cuerpo reemplazado contendrá operaciones distintas de la operación fusionada. 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, en el que solo existiría la operación fusionada.
A continuación, se incluye 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í se muestra un fragmento de código que asigna esta operación compuesta a una operación fusionada en LiteRT 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());
}