5.5 TensorFlow.js Web 端部署

文档摘要

5.5 TensorFlow.js Web 端部署 TensorFlow.js Web 端部署详解 TensorFlow.js Web 端部署方式 TensorFlow.js 提供了多种在 Web 端部署模型的方式，主要包括：加载预训练模型：这是最常见的场景，即在 TensorFlow 或 Keras 中训练好模型，然后将其转换为 TensorFlow.js 格式，并在 Web 应用程序中加载。直接在浏览器中训练模型：虽然不常见，但 TensorFlow.js 也支持在浏览器中直接训练模型。这对于在线学习或数据隐私敏感的场景非常有用。使用预训练模型进行迁移学习：利用现有的预训练模型作为基础，在其上进行微调，以适应特定的任务。加载预训练模型这是最常见的部署方式。

5.5 TensorFlow.js Web 端部署

TensorFlow.js Web 端部署详解

1. TensorFlow.js Web 端部署方式

TensorFlow.js 提供了多种在 Web 端部署模型的方式，主要包括：

加载预训练模型： 这是最常见的场景，即在 TensorFlow 或 Keras 中训练好模型，然后将其转换为 TensorFlow.js 格式，并在 Web 应用程序中加载。
直接在浏览器中训练模型： 虽然不常见，但 TensorFlow.js 也支持在浏览器中直接训练模型。这对于在线学习或数据隐私敏感的场景非常有用。
使用预训练模型进行迁移学习： 利用现有的预训练模型作为基础，在其上进行微调，以适应特定的任务。

2. 加载预训练模型

这是最常见的部署方式。通常，我们使用 Python 中的 TensorFlow 或 Keras 训练模型，然后将其转换为 TensorFlow.js 格式。

2.1 模型转换

首先，需要将 TensorFlow 或 Keras 模型转换为 TensorFlow.js 可以识别的格式。 TensorFlow.js 提供了 tensorflowjs_converter 工具来完成这项工作。


pip install tensorflowjs

假设我们有一个名为 my_model.h5 的 Keras 模型，可以使用以下命令将其转换为 TensorFlow.js 格式：


tensorflowjs_converter --input_format keras my_model.h5 tfjs_model

这将在 tfjs_model 目录下生成两个文件：model.json（模型结构）和 weights.bin（模型权重）。

2.2 Web 端加载模型

在 HTML 文件中，引入 TensorFlow.js 库：


<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>

然后，使用 tf.loadLayersModel() 函数加载模型：


async function loadModel() {
  try {
    const model = await tf.loadLayersModel('tfjs_model/model.json');
    console.log('模型加载成功');
    return model;
  } catch (error) {
    console.error('模型加载失败:', error);
    return null;
  }
}
let model;
loadModel().then(loadedModel => {
  if (loadedModel) {
    model = loadedModel;
    // 模型加载成功后，可以进行预测
    // 例如：predict(model, inputData);
  }
});

2.3 模型预测

加载模型后，就可以使用 model.predict() 函数进行预测。


async function predict(model, inputData) {
  tf.engine().startScope(); // 开启内存管理作用域
  try {
    // 将输入数据转换为 TensorFlow.js 张量
    const tensor = tf.tensor(inputData, [1, inputData.length]); // 假设输入数据是一维数组
    // 进行预测
    const prediction = model.predict(tensor);
    // 获取预测结果
    const result = await prediction.data();
    console.log('预测结果:', result);
    tf.engine().endScope(); // 结束内存管理作用域
    return result;
  } catch (error) {
    console.error('预测出错:', error);
    tf.engine().endScope(); // 确保即使出错也结束作用域
    return null;
  } finally {
    // 可选：清理张量，释放内存
    // tensor.dispose();
    // prediction.dispose();
  }
}

2.4 完整示例


<!DOCTYPE html>
<html>
<head>
  <title>TensorFlow.js 模型加载示例</title>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>
</head>
<body>
  <h1>TensorFlow.js 模型加载示例</h1>
  <script>
    async function loadModel() {
      try {
        const model = await tf.loadLayersModel('tfjs_model/model.json');
        console.log('模型加载成功');
        return model;
      } catch (error) {
        console.error('模型加载失败:', error);
        return null;
      }
    }
    async function predict(model, inputData) {
        tf.engine().startScope(); // 开启内存管理作用域
        try {
            // 将输入数据转换为 TensorFlow.js 张量
            const tensor = tf.tensor(inputData, [1, inputData.length]); // 假设输入数据是一维数组
            // 进行预测
            const prediction = model.predict(tensor);
            // 获取预测结果
            const result = await prediction.data();
            console.log('预测结果:', result);
            tf.engine().endScope(); // 结束内存管理作用域
            return result;
        } catch (error) {
            console.error('预测出错:', error);
            tf.engine().endScope(); // 确保即使出错也结束作用域
            return null;
        } finally {
            // 可选：清理张量，释放内存
            // tensor.dispose();
            // prediction.dispose();
        }
    }
    let model;
    loadModel().then(loadedModel => {
      if (loadedModel) {
        model = loadedModel;
        // 模型加载成功后，可以进行预测
        const inputData = [0.1, 0.2, 0.3, 0.4]; // 示例输入数据
        predict(model, inputData);
      }
    });
  </script>
</body>
</html>

Mermaid 图示：模型加载与预测流程

3. 直接在浏览器中训练模型

虽然不如加载预训练模型常见，但 TensorFlow.js 也支持在浏览器中训练模型。这对于一些特定的场景非常有用，例如：

在线学习： 模型可以根据用户的实时数据进行调整。
数据隐私： 数据保留在客户端，无需发送到服务器进行训练。

3.1 创建模型

首先，需要定义模型的结构。可以使用 tf.sequential() 创建一个顺序模型，或者使用 tf.model() 创建一个更复杂的模型。


const model = tf.sequential();
model.add(tf.layers.dense({units: 1, activation: 'linear', inputShape: [1]}));
// 定义优化器和损失函数
model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

3.2 准备训练数据

准备训练数据，包括输入数据 (xs) 和目标数据 (ys)。


const xs = tf.tensor2d([1, 2, 3, 4], [4, 1]);
const ys = tf.tensor2d([2, 4, 6, 8], [4, 1]);

3.3 训练模型

使用 model.fit() 函数训练模型。


async function trainModel(model, xs, ys) {
  await model.fit(xs, ys, {epochs: 100});
  console.log('模型训练完成');
}
trainModel(model, xs, ys);

3.4 模型预测

训练完成后，可以使用 model.predict() 函数进行预测。


const input = tf.tensor2d([5], [1, 1]);
const prediction = model.predict(input);
const result = prediction.dataSync()[0];
console.log('预测结果:', result);

3.5 完整示例


<!DOCTYPE html>
<html>
<head>
  <title>TensorFlow.js 浏览器端训练示例</title>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>
</head>
<body>
  <h1>TensorFlow.js 浏览器端训练示例</h1>
  <script>
    // 定义模型
    const model = tf.sequential();
    model.add(tf.layers.dense({units: 1, activation: 'linear', inputShape: [1]}));
    // 定义优化器和损失函数
    model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});
    // 准备训练数据
    const xs = tf.tensor2d([1, 2, 3, 4], [4, 1]);
    const ys = tf.tensor2d([2, 4, 6, 8], [4, 1]);
    // 训练模型
    async function trainModel(model, xs, ys) {
      await model.fit(xs, ys, {epochs: 100});
      console.log('模型训练完成');
      // 预测
      const input = tf.tensor2d([5], [1, 1]);
      const prediction = model.predict(input);
      const result = prediction.dataSync()[0];
      console.log('预测结果:', result);
    }
    trainModel(model, xs, ys);
  </script>
</body>
</html>

Mermaid 图示：浏览器端训练流程

4. 使用预训练模型进行迁移学习

迁移学习是一种利用预训练模型作为基础，在其上进行微调，以适应特定任务的技术。这可以显著减少训练时间和数据需求。

4.1 加载预训练模型

首先，加载一个预训练模型，例如 MobileNet 或 VGG19。


async function loadPretrainedModel() {
  try {
    const model = await tf.loadLayersModel('https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v2_100_224/classification/5/model.json');
    console.log('预训练模型加载成功');
    return model;
  } catch (error) {
    console.error('预训练模型加载失败:', error);
    return null;
  }
}
let pretrainedModel;
loadPretrainedModel().then(loadedModel => {
  if (loadedModel) {
    pretrainedModel = loadedModel;
    // 冻结预训练模型的权重，防止在微调过程中被修改
    for (let i = 0; i < pretrainedModel.layers.length; i++) {
      pretrainedModel.layers[i].trainable = false;
    }
    // 在预训练模型的基础上添加新的层
    const newModel = tf.sequential();
    newModel.add(pretrainedModel);
    newModel.add(tf.layers.dense({units: 10, activation: 'relu'}));
    newModel.add(tf.layers.dense({units: 1, activation: 'sigmoid'})); // 根据你的任务修改激活函数
    // 编译模型
    newModel.compile({loss: 'binaryCrossentropy', optimizer: 'adam', metrics: ['accuracy']});
    // 现在可以使用你的数据训练 newModel
    // 例如：trainNewModel(newModel, trainingData, trainingLabels);
  }
});

4.2 添加自定义层

在预训练模型的基础上，添加自定义的层，以适应特定的任务。


// 在预训练模型的基础上添加新的层
const newModel = tf.sequential();
newModel.add(pretrainedModel);
newModel.add(tf.layers.dense({units: 10, activation: 'relu'}));
newModel.add(tf.layers.dense({units: 1, activation: 'sigmoid'})); // 根据你的任务修改激活函数

4.3 编译和训练模型

编译模型，并使用自己的数据进行训练。


// 编译模型
newModel.compile({loss: 'binaryCrossentropy', optimizer: 'adam', metrics: ['accuracy']});
// 训练模型
async function trainNewModel(model, trainingData, trainingLabels) {
  await model.fit(trainingData, trainingLabels, {epochs: 10});
  console.log('迁移学习模型训练完成');
}
// 假设 trainingData 和 trainingLabels 已经准备好
// trainNewModel(newModel, trainingData, trainingLabels);

Mermaid 图示：迁移学习流程

5. 性能优化

在 Web 端部署 TensorFlow.js 模型时，性能优化至关重要。以下是一些常用的优化技巧：

使用 WebGL 后端： WebGL 后端利用 GPU 加速计算，可以显著提高性能。确保浏览器支持 WebGL，并使用 tf.setBackend('webgl') 启用它。
量化模型： 量化可以将模型的权重从 32 位浮点数转换为 8 位整数，从而减小模型大小并提高推理速度。 TensorFlow.js 提供了量化 API。
延迟加载模型： 只在需要时才加载模型，可以减少初始加载时间。
使用异步操作： 避免阻塞主线程，使用 async/await 或 Promise 处理耗时操作。
管理内存： TensorFlow.js 使用 WebGL 后端时，需要手动管理内存。使用 tf.dispose() 释放不再需要的张量。也可以使用 tf.tidy() 函数自动清理临时张量。

6. 总结

TensorFlow.js 为 Web 端 AI 应用开发提供了强大的工具。通过加载预训练模型、直接在浏览器中训练模型或使用迁移学习，我们可以构建各种各样的智能应用。性能优化是 Web 端部署的关键，需要根据实际情况选择合适的优化策略。通过本文的介绍，希望能帮助您更好地理解和应用 TensorFlow.js 在 Web 端部署。

作者与出处

原作者：灏天文库

来源：灏天文库

原文： https://aiknowledge.cn/collection/219-tensorflow%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0%E5%AE%9E%E6%88%98%E6%95%99%E7%A8%8B

整理：灏天文库整理

由灏天文库平台收录，内容或由平台用户上传，仅供学习交流