概览
本页面介绍了转换复合操作所需的设计和步骤 转换为 LiteRT 中的融合操作。这种基础架构 通用，并且支持转换 TensorFlow 中的任何复合操作 LiteRT 中相应的融合操作。
此基础架构的一个示例用途是将 TensorFlow RNN 运算融合到 LiteRT（详见此处）。
什么是融合操作
TensorFlow 操作可以是基元操作，例如 tf.add 或者它们可以是 由其他原初操作（如 tf.einsum。基元 操作在 TensorFlow 图中显示为单个节点，而复合 操作是 TensorFlow 图中的节点集合。执行 复合操作等同于执行其每个组成部分原语 操作。
融合操作对应于单个操作，该操作将 由每个原初操作执行的 复合操作。
融合操作的优势
融合操作旨在最大限度地提高底层内核的性能 优化整体计算并减少内存用量， 足迹。这非常有用，尤其是对于低延迟的推理工作负载 以及资源受限的移动平台
融合操作还提供更高级别的接口来定义复杂的 像量化这样的转换，否则，这些转换是不可行的， 难以在更精细的级别上执行。
由于 如上文所述。这些融合操作通常对应于复合 操作。中的复合操作示例 在 LiteRT 中作为单个融合操作实现的 TensorFlow 包括各种 RNN 运算，例如单向和双向序列 LSTM、卷积（conv2d、bias add、relu）、全连接（matmul、biasly add、 relu）等。在 LiteRT 中，LSTM 量化目前仅 混合 LSTM 操作中实现的。
融合操作带来的挑战
将复合操作从 TensorFlow 转换为融合操作 LiteRT 是个难题。原因如下：
复合操作在 TensorFlow 图中表示为一组 原始操作没有明确定义的边界。结果可能非常 很难识别（例如，通过模式匹配）子图 。
可能有多个针对融合 LiteRT 操作。例如，许多 LSTM 实现 （Keras、Babelfish/lingvo 等）， 不同的原始操作，但它们仍然可以转换为 LiteRT 中相同的融合 LSTM 操作。
因此，事实证明，转换融合操作极具挑战性。
从复合操作转换为 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 解释器实现自定义操作和注册
将融合操作实现为 TFLite 自定义操作 - 请参阅 说明。
请注意，用于注册操作的名称应类似于
在实现签名时，在 name 属性中指定此属性。
示例中的操作示例是
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 函数注解。请参阅嵌入查询的示例。通过 函数定义了接口，并且应使用其参数来实现 转换逻辑。
写入转化代码
根据函数的接口编写转化代码，
implements 注解。查看嵌入的融合示例
lookup。从概念上讲，转化代码会替换
该接口的实现。
在 prepare-composite-functions 环节中，转换中的插件 代码。
在更高级的用法中，可以对 复合运算的运算数，以便推导融合运算数 操作。请参阅 Keras LSTM. 以转换代码为例
转换为 LiteRT
使用 TFLiteConverter.from_saved_model API 转换为 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
- 融合复合运算
- 变量可变性分析
- 冻结所有只读变量
这样，我们就可以使用表示 在函数内联和变量冻结之前的复合操作。
实现操作融合
我们来更详细地了解一下 Fusion Pass 操作。此卡券会 以下：
- 循环遍历 MLIR 模块中的所有函数。
- 如果函数具有 tf._implements 属性（根据该属性 值，调用相应的操作融合实用程序。
- 运算融合实用程序对函数的运算数进行运算， 属性（用作转化的接口）并替换 函数正文与包含 融合操作。
- 在许多情况下，替换后的正文会包含 融合操作。这些转换对应于 函数的操作数，以获取融合运算的运算数。 由于这些计算都可以是恒定折叠式的， 存在于导出的 Flatbuffer 中，其中只有融合操作 存在。
以下是卡券中的代码段，其中显示了主要工作流：
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 */
}
以下代码段展示了如何将此复合操作映射到 Fused 操作作为转换接口。
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());
}