Descripción general
En esta página, se describen el diseño y los pasos necesarios para convertir operaciones compuestas en TensorFlow con operaciones fusionadas en LiteRT. Esta infraestructura está de uso general y admite la conversión de cualquier operación compuesta en TensorFlow a la operación fusionada correspondiente en LiteRT.
Un ejemplo de uso de esta infraestructura es TensorFlow RNN Operations Fusion para LiteRT, como se detalla aquí.
¿Qué son las operaciones fusionadas?
Las operaciones de TensorFlow pueden ser operaciones primitivas, p.ej., tf.add o pueden ser compuesta por otras operaciones primitivas, por ejemplo, tf.einsum. Un componente primitivo se muestra como un solo nodo en el grafo de TensorFlow, mientras que una es una colección de nodos en el grafo de TensorFlow. Ejecutar un una operación compuesta equivale a ejecutar cada una de las operaciones primitivas que la constituyen. las operaciones.
Una operación fusionada corresponde a una única operación que suma todos los procesamiento que realiza cada operación primitiva en el rango una operación compuesta.
Beneficios de las operaciones fusionadas
Existen operaciones fusionadas para maximizar el rendimiento de su kernel subyacente a través de la optimización del procesamiento general y la reducción de la memoria para reducir su huella de seguridad. 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 modelos como la cuantización, que, de otro modo, sería inviable o sería muy difíciles de hacer a un nivel más detallado.
LiteRT tiene muchas instancias de operaciones fusionadas por estos motivos: que se explicó anteriormente. Por lo general, estas operaciones fusionadas corresponden a la estructura las operaciones en el programa de origen de TensorFlow. Ejemplos de operaciones compuestas en TensorFlow que se implementan como una única operación fusionada en LiteRT Incluyen varias operaciones de RNN, como secuencias unidireccionales y bidireccionales. LSTM, convolución (conv2d, sesgo agregado, relu), completamente conectado (matmul, sesgo agregado, relu) y mucho más. En LiteRT, la cuantización de LSTM solo está disponible actualmente. implementadas en las operaciones fusionadas de LSTM.
Desafíos de las operaciones fusionadas
Convertir operaciones compuestas de TensorFlow en operaciones fusionadas 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 sin un límite bien definido. Puede ser muy difícil identificar (por ejemplo, a través de coincidencia de patrones) el subgrafo correspondiente a esa operación compuesta.
Puede haber más de una implementación de TensorFlow dirigida a una fusión LiteRT. Por ejemplo, existen muchas implementaciones de LSTM en TensorFlow (Keras, Babelfish/lingvo, etc.), y cada uno de ellos se compone de operaciones primitivas diferentes, pero todas aún podrían convertirse la misma operación de LSTM fusionada en LiteRT.
Por ello, la conversión de operaciones fusionadas ha demostrado ser todo un desafío.
Pasa de una operación compuesta a una operación personalizada de TFLite (recomendado)
Une la operación compuesta en una tf.function
En muchos casos, parte del modelo
puede asignarse a una única operación
TFLite. Esto puede mejorar el rendimiento a la hora de escribir una implementación optimizada.
para operaciones específicas. Para poder crear una operación fusionada en TFLite,
identificar la parte del gráfico que representa una operación fusionada y unirla
una tf.function con
“experimental_implements” a un tf.function, que tiene atributos
el valor tfl_fusable_op con el valor true. Si la operación personalizada tarda
atributos y, luego, los pasan como parte de la misma "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 configures allow_custom_ops en el conversor como
El atributo tfl_fusable_op implica que ya lo es.
Implementa operaciones personalizadas y regístrate con el intérprete de TFLite
Implementa tu operación fusionada como una operación personalizada de TFLite. Consulta instrucciones.
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 implementación de 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, ®);
Convertir de una operación compuesta a una fusionada (avanzado)
Arquitectura general para convertir las operaciones compuestas de TensorFlow en A continuación, se muestran las operaciones combinadas de LiteRT:
Une la operación compuesta en una tf.function
En el código fuente del modelo de TensorFlow, identifica y abstrae la composición una operación en una tf.function con el experimental_implements la anotación de función. Mira un ejemplo de una búsqueda de incorporación. El define la interfaz, y sus argumentos deberían usarse para implementar la lógica de conversión.
Escribir el código de conversión
El código de conversión se escribe según la interfaz de la función con el
implements. Ver un ejemplo de fusión para la incorporación
búsqueda. Conceptualmente, el código de conversión reemplaza al compuesto
implementación de esta interfaz con la fusionada.
En el pase prepare-composite-functions, complemente en su conversión automático.
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 una sola operación. Ver Keras LSTM. el código de conversión como ejemplo.
Convertir a LiteRT
Usa el TFLiteConverter.from_saved_model para convertir a LiteRT.
Detrás de escena
Ahora describimos los detalles de alto nivel del diseño general en la conversión a fusión operaciones en LiteRT.
Cómo redactar operaciones en TensorFlow
El uso de tf.function con el experimental_implements de la función permite a los usuarios redactar de forma explícita operaciones nuevas con las operaciones primitivas de TensorFlow y especificar la interfaz a la que de cada operación compuesta. Esto es muy útil, ya que proporciona lo siguiente:
- Un límite bien definido para la operación compuesta en la Grafo de TensorFlow.
- Especifica de forma explícita la interfaz que implementa esta operación. El argumentos del tf.function corresponden a los argumentos de esta interfaz.
A modo de ejemplo, consideremos una operación compuesta definida para implementar búsqueda de incorporaciones. Esto se relaciona con 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
Haciendo que los modelos usen operaciones compuestas a través de tf.function como ilustradas anteriormente, es posible crear una infraestructura general para identificar y convertir esas operaciones en operaciones fusionadas de LiteRT.
Cómo extender el conversor de LiteRT
El conversor LiteRT que se lanzó a principios de este año solo era compatible importar modelos de TensorFlow como un grafo en el que todas las variables se reemplazan por sus los valores constantes correspondientes. Esto no funciona para la fusión de operaciones, esos gráficos tienen todas las funciones intercaladas para que las variables se puedan convertir en constantes.
Para aprovechar el
tf.function con el
experimental_implements durante el proceso de conversión, las funciones
y deben conservarse hasta más adelante en el proceso de conversión.
Por ello, implementamos un nuevo flujo de trabajo para importar y convertir TensorFlow. en el conversor para admitir el caso de uso de fusión de operaciones compuestas. Específicamente, estas son las funciones nuevas que se agregaron:
- Importar modelos guardados TensorFlow a MLIR
- fusionar operaciones compuestas
- análisis de mutabilidad de variables
- inmovilizar todas las variables de solo lectura
Esto nos permite llevar a cabo la fusión de operaciones usando las funciones que representan la operaciones compuestas antes del intercalado de funciones y la inmovilización de variables.
Implementa la fusión de operaciones
Veamos con más detalle el pase de fusión de operaciones. Este pase hace lo siguiente:
- Repetir todas las funciones en el módulo de MLIR
- Si una función tiene el atributo tf._implements, basado en el atributo de salida, llama a la utilidad de fusión de operaciones adecuada.
- La utilidad de fusión de operaciones opera en los operandos de la función y atributos (que sirven como interfaz para la conversión) y reemplaza el cuerpo de la función con un cuerpo de función equivalente que contenga las una operación fusionada.
- En muchos casos, el cuerpo reemplazado contendrá operaciones distintas de una operación fusionada. Estas corresponden a algunas transformaciones estáticas operandos de la función para obtener los operandos de la operación fusionada. Dado que todos estos cálculos pueden ir constantes, no serían presente en el búfer plano exportado en el que solo la operación fusionada existen.
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 fusión en LiteRT que aprovecha la función como 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());
}