使用 Keras 和 TensorFlow Lite 的设备端大型语言模型

将大型语言模型部署到Android设备上是一个复杂的任务，因为这些模型通常需要大量的计算资源。首先，您需要加载和优化模型，然后选择一个合适的推理引擎，最后将其集成到Android应用程序中。以下是一般步骤：

1. 加载和优化模型：

• 如果您使用Keras构建了一个大型语言模型，首先需要将模型加载到Python中。
• 使用TensorFlow的转换工具（如TensorFlow Lite Converter）将模型转换为TensorFlow Lite格式，这可以减小模型的大小并提高在移动设备上的性能。
• 还可以使用量化技术将模型压缩，以减小内存和计算要求。

2. 选择推理引擎：

• 在Android设备上运行模型需要选择一个合适的推理引擎。TensorFlow Lite和ONNX Runtime是常见的选择，可以在Android上运行各种类型的模型。
• 评估不同引擎的性能，并选择最适合您模型和应用需求的引擎。

3. 集成到Android应用程序中：

• 创建一个Android应用项目，并将模型文件和推理引擎集成到项目中。
• 使用Android Studio开发应用界面，以便用户可以与模型进行交互。
• 在应用中设置推理管道，以处理用户输入并获取模型的输出。

4. 优化性能：

• 在移动设备上运行大型语言模型可能会导致性能问题，特别是内存和计算需求。
• 优化模型推理，例如批量处理输入，减少内存占用，使用GPU加速等。
• 考虑使用模型缓存，以避免重复加载模型。

5. 测试和调试：

• 在实际设备上进行测试，确保模型在移动设备上表现良好。
• 进行性能分析和调试，以解决任何性能或功能问题。

6. 安全性和隐私：

• 在应用程序中实施必要的安全性和隐私措施，特别是处理敏感数据时。
• 确保模型和数据的安全存储和传输。

7. 发布应用程序：

• 准备应用程序以发布到Google Play商店或其他适当的应用商店。

请注意，这是一个高级任务，需要深入的移动应用程序开发和深度学习知识。您可能需要详细了解Android开发、TensorFlow Lite或ONNX Runtime，以及模型优化技术，以成功地将大型语言模型部署到Android设备上。同时，确保遵守相关的法规和隐私政策，以保护用户数据。

相关示例

近期，最令人兴奋的机器学习突破之一是大型语言模型 (LLM)。这些模型可用于生成文字、翻译语言以及以全面且信息丰富的方式回答问题。LLM（如 Google LaMDA 和 PaLM）已基于大量文本数据进行训练，使其能够学习字词和短语之间的统计模式和关系。这样，它们就可以生成类似于人工撰写的文本，并能够很准确地进行翻译。

LLM 会占用很多存储空间，并且通常需要消耗大量的算力才能运行，这意味着它们通常部署在云端，并且由于移动设备上的计算能力有限，对设备端机器学习 (ODML) 而言极具挑战性。不过，您也可以在新型 Android 设备上运行较小规模的 LLM（例如 GPT-2），并且仍能取得令人瞩目的成果。

在此 Codelab 中，您将学习如何通过以下方式构建由 LLM 提供支持的应用：

• 使用 KerasNLP 加载预训练的 LLM (GPT-2)
• 使用 KerasNLP 对 LLM (GPT-2) 进行微调
• 使用 TensorFlow Lite 在 Android 上转换、优化和部署 LLM

前提条件

• 掌握 Keras 和 TensorFlow Lite 的中级知识
• 掌握 Android 开发方面的基础知识

学习内容

• 如何使用 KerasNLP 加载和微调预先训练的 LLM
• 如何量化 LLM 并将其转换为 TensorFlow Lite
• 如何使用转换后的 TensorFlow Lite 模型进行推断

关于 KerasNLP

KerasNLP 是一个自然语言处理库，可为用户的整个开发周期提供支持。我们的工作流程由模块化组件构建，这些组件在开箱即用时具有最先进的预设权重和架构，并且在需要更多控制时可轻松定制。我们强调所有工作流程的图内计算，以便开发人员可以使用 TensorFlow 生态系统轻松实现生产。该库是核心 Keras API 的扩展；所有高级模块都 Layers接受 Models与核心 Keras 相同水平的打磨。

• 安装环境

pip install git+https://github.com/keras-team/keras-nlp.git --upgrade

• 快速开始

使用 API 对小型情感分析任务微调 BERT keras_nlp.models：

import keras_nlp
import tensorflow_datasets as tfds

imdb_train, imdb_test = tfds.load(
    "imdb_reviews",
    split=["train", "test"],
    as_supervised=True,
    batch_size=16,
)
# Load a BERT model.
classifier = keras_nlp.models.BertClassifier.from_preset("bert_base_en_uncased")
# Fine-tune on IMDb movie reviews.
classifier.fit(imdb_train, validation_data=imdb_test)
# Predict two new examples.
classifier.predict(["What an amazing movie!", "A total waste of my time."])

将 LLM 部署到设备端

• Colab

• TensorFlow Codelab GitHub

Android 应用使用

• Github :examples/lite/examples/generative_ai

初始化模型
initModel 方法用于初始化 TensorFlow-Lite 模型。它首先尝试加载模型文件，然后实例化一个 Interpreter 对象来运行模型。

override suspend fun initModel(): InitModelResult {
    
    
    return withContext(dispatcher) {
    
    
        // Load model file
        val loadResult = loadModelFile(context)
        // Determine if load was successful
        if (loadResult.isFailure) {
    
    
            val exc = loadResult.exceptionOrNull()
            return@withContext if (exc is FileNotFoundException) {
    
    
                InitModelResult.Error(AutoCompleteServiceError.MODEL_FILE_NOT_FOUND)
            } else {
    
    
                InitModelResult.Error(AutoCompleteServiceError.MODEL_NOT_INITIALIZED)
            }
        }
        // Instantiate interpreter with loaded model
        val model = loadResult.getOrNull()
        isInitialized = model?.let {
    
    
            interpreter = Interpreter(it)
            true
        } ?: false
        if (isInitialized) InitModelResult.Success
        else InitModelResult.Error(AutoCompleteServiceError.MODEL_NOT_INITIALIZED)
    }
}

运行模型

runInterpreterOn 方法用于运行 TensorFlow-Lite 模型，根据输入文本生成新的文本。

@WorkerThread
private fun runInterpreterOn(input: String): String {
    
    
    outputBuffer.clear()
    // Run interpreter, which will generate text into outputBuffer
    interpreter.run(input, outputBuffer)
    // Set output buffer limit to current position & position to 0
    outputBuffer.flip()
    // Get bytes from output buffer
    val bytes = ByteArray(outputBuffer.remaining())
    outputBuffer.get(bytes)
    outputBuffer.clear()
    // Return bytes converted to String
    return String(bytes, Charsets.UTF_8)
}

使用指南

步骤 1. 使用 Keras 训练语言模型
在本次演示中，我们将使用 KerasNLP 来获取 GPT-2 模型。KerasNLP 是一个库，包含用于自然语言处理任务的最先进的预训练模型，并且可以在用户的​​整个开发周期中为用户提供支持。您可以在KerasNLP 存储库中查看可用模型的列表。工作流程由模块化组件构建，这些组件在开箱即用时具有最先进的预设权重和架构，并且在需要更多控制时可轻松定制。创建 GPT-2 模型可以通过以下步骤完成：

gpt2_tokenizer = keras_nlp.models.GPT2Tokenizer.from_preset("gpt2_base_en")

gpt2_preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
  "gpt2_base_en",
  sequence_length=256,
  add_end_token=True,
)

gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
  "gpt2_base_en", 
  preprocessor=gpt2_preprocessor,
)

您可以在 GitHub 上查看完整的 GPT-2 模型实现。

步骤 2. 将 Keras 模型转换为 TFLite 模型
generate()从执行转换的 GPT2CausalLM 函数开始。包装generate()函数以创建具体的 TensorFlow 函数：

@tf.function
def generate(prompt, max_length):
  # prompt: input prompt to the LLM in string format
  # max_length: the max length of the generated tokens 
  return gpt2_lm.generate(prompt, max_length)
concrete_func = generate.get_concrete_function(tf.TensorSpec([], tf.string), 100)

现在定义一个辅助函数，它将使用输入和 TFLite 模型运行推理。TensorFlow 文本操作不是 TFLite 运行时中的内置操作，因此您需要添加这些自定义操作，以便解释器对此模型进行推理。该辅助函数接受输入和执行转换的函数，即generator()上面定义的函数。

def run_inference(input, generate_tflite):
  interp = interpreter.InterpreterWithCustomOps(
    model_content=generate_tflite,
    custom_op_registerers=tf_text.tflite_registrar.SELECT_TFTEXT_OPS)
  interp.get_signature_list()

  generator = interp.get_signature_runner('serving_default')
  output = generator(prompt=np.array([input]))

您现在可以转换模型：

gpt2_lm.jit_compile = False
converter = tf.lite.TFLiteConverter.from_concrete_functions(
  [concrete_func],
  gpt2_lm)

converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TFLite ops
  tf.lite.OpsSet.SELECT_TF_OPS, # enable TF ops
]
converter.allow_custom_ops = True
converter.target_spec.experimental_select_user_tf_ops = [
  "UnsortedSegmentJoin",
  "UpperBound"
]
converter._experimental_guarantee_all_funcs_one_use = True
generate_tflite = converter.convert()
run_inference("I'm enjoying a", generate_tflite)

步骤 3. 量化
TensorFlow Lite 实现了一种称为量化的优化技术，可以减小模型大小并加速推理。通过量化过程，32 位浮点数被映射为更小的 8 位整数，因此将模型大小减少了 4 倍，以便在现代硬件上更有效地执行。在 TensorFlow 中进行量化有多种方法。您可以访问TFLite 模型优化和TensorFlow 模型优化工具包页面以获取更多信息。下面简要解释量化的类型。

在这里，您将通过将转换器优化标志设置为 tf.lite.Optimize.DEFAULT，在 GPT-2 模型上使用训练后动态范围量化，其余转换过程与之前详述的相同。我们测试发现，使用这种量化技术，最大输出长度设置为 100 时，Pixel 7 上的延迟约为 6.7 秒。

gpt2_lm.jit_compile = False
converter = tf.lite.TFLiteConverter.from_concrete_functions(
  [concrete_func],
  gpt2_lm)

converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TFLite ops
  tf.lite.OpsSet.SELECT_TF_OPS, # enable TF ops
]
converter.allow_custom_ops = True
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.experimental_select_user_tf_ops = [
  "UnsortedSegmentJoin",
  "UpperBound"
]
converter._experimental_guarantee_all_funcs_one_use = True
quant_generate_tflite = converter.convert()
run_inference("I'm enjoying a", quant_generate_tflite)

with open('quantized_gpt2.tflite', 'wb') as f:
  f.write(quant_generate_tflite)

步骤 4. Android 应用程序集成
您可以克隆此存储库并替换android/app/src/main/assets/autocomplete.tflite为转换后的quant_generate_tflite文件。

相关参考

使用 Keras 和 TensorFlow Lite 的设备端大型语言模型
https://codelabs.developers.google.com/kerasnlp-tflite