TensorFlow.js:将 Python SavedModel 转换为 TensorFlow.js 格式

1. 简介

您已经开始使用 TensorFlow.js,试用了我们的预建模型,甚至可能还创建了自己的模型。不过,您看到 Python 中出现了一些前沿研究,并且好奇这些研究是否能在网络浏览器中运行,从而以可扩缩的方式让数百万人实现您想到的那个绝妙创意。听起来是不是很熟悉?如果是这样,那么本 Codelab 正适合您!

TensorFlow.js 团队开发了一款便捷的工具,可通过命令行转换器将 SavedModel 格式的模型转换为 TensorFlow.js 模型,这样您就可以在网络上使用这些模型,并充分发挥其覆盖面广、可大规模应用的优势。

学习内容

在此 Codelab 中,您将学习如何使用 TensorFlow.js 命令行转换器将 Python 生成的 SavedModel 移植到 model.json 格式,以便在 Web 浏览器中在客户端执行。

具体而言:

  • 如何创建简单的 Python 机器学习模型并将其保存为 TensorFlow.js 转换器所需的格式。
  • 如何安装和使用 TensorFlow.js 转换器来处理从 Python 导出的 SavedModel。
  • 获取转换后生成的文件,并在您的 JS Web 应用中使用。
  • 了解在出现问题(并非所有模型都会转换)时该怎么做,以及您有哪些选择。

想象一下,您能够利用新发布的研究成果,让全球数百万 JavaScript 开发者使用该模型。或者,您也可以在自己的作品中使用此功能,这样一来,如果您的作品在网络浏览器中运行,那么世界上的任何人都可以体验,因为无需复杂的依赖项或环境设置。准备好开始黑客马拉松了吗?现在就试试吧!

与我们分享您的转化成果!

您可以运用今天学到的知识,尝试将一些自己喜爱的模型从 Python 进行转换。如果您成功完成了上述操作,并制作了一个展示模型实际应用的演示网站,请在社交媒体上使用 #MadeWithTFJS 标签标记我们,这样就有机会在我们的 TensorFlow 博客甚至未来的展示与推介活动中展示您的项目。我们希望看到更多出色的研究成果被移植到 Web 上,让更多人能够以创新或富有创意的方式使用此类模型,就像这个出色的示例一样。

2. 什么是 TensorFlow.js?

1aee0ede85885520.png

TensorFlow.js 是一个开源机器学习库,可在任何能够运行 JavaScript 的环境中运行。它以用 Python 编写的原始 TensorFlow 库为基础,旨在面向 JavaScript 生态系统打造全新的开发者体验以及 API 体验。

它的应用场合如何?

鉴于 JavaScript 的可移植性,您现在可以使用一种语言编写机器学习代码,然后轻松地在下列所有平台上运行此代码:

  • 使用 Vanilla JavaScript 的客户端网络浏览器
  • 使用 Node.js 的服务器端甚至是 Raspberry Pi 等 IoT 设备
  • 使用 Electron 的桌面应用
  • 使用 React Native 的原生移动应用

TensorFlow.js 还支持在这些环境中使用多个后端(这些环境是它可以在其中运行的实际硬件环境,例如 CPU 或 WebGL 等。此处的“后端”并不是指服务器端环境,例如,用于执行的后端可能是 WebGL 中的客户端),以确保兼容性,同时保持快速运行。TensorFlow.js 目前支持:

  • 在设备显卡 (GPU) 上运行 WebGL - 这是使用 GPU 加速运行较大模型(大小超过 3MB)的最快方式。
  • 在 CPU 上执行 Web Assembly (WASM) - 旨在提升包括老款手机在内的设备的 CPU 性能。这种方式更适合较小的模型(大小在 3MB 以下),由于将内容上传到图形处理器存在开销,因此在 CPU 上运行 WASM 可能实际上比运行 WebGL 要快。
  • CPU 执行 - 如果其他环境都不可用,则回退至此方式。这是三种方式中运行最慢的一种,但始终可以使用。

注意:如果您知道要在哪个设备上运行,则可以选择强制调用这些后端之一。或者,如果您没有指定,可以直接让 TensorFlow.js 为您做决定。

客户端的强大功能

在客户端计算机的网络浏览器中运行 TensorFlow.js 具有诸多值得关注的优势。

隐私权

您可以在客户端计算机上对数据进行训练和分类,而不将数据发送到第三方网络服务器。有时候,这可能是一项强制要求,例如 GDPR 等地方法律要求这么做,或者对于您所处理的数据,用户可能想要将其保留在自己的计算机上,而不发送给第三方。

速度

由于您无需向远程服务器发送数据,因此推理(一种对数据进行分类的操作)可能会更快。更棒的是,如果用户向您授予权限,您可以直接访问设备的传感器,比如相机、麦克风、GPS、加速度计等。

覆盖面和规模

只要点击一下您发送的链接,世界上的任何人就可以在他们的浏览器中打开相应网页,并利用您创作的内容。无需使用 CUDA 驱动程序等进行复杂的服务器端 Linux 设置,即可使用该机器学习系统。

费用

没有服务器意味着您只需为托管 HTML、CSS、JS 和模型文件的内容分发网络 (CDN) 付费。CDN 的费用要比全天候运行服务器(可能搭载了显卡)便宜得多。

服务器端功能

利用 TensorFlow.js 的 Node.js 实现可获得以下特性。

全面的 CUDA 支持

在服务器端,要实现显卡加速,必须安装 NVIDIA CUDA 驱动程序,以便 TensorFlow 能够使用显卡(与使用 WebGL 的浏览器不同,这种情况下无需安装该驱动程序)。不过,借助全面的 CUDA 支持,您可以充分利用显卡更精细的功能,从而缩短训练和推理时间。性能与 Python TensorFlow 实现相当,因为它们具有相同的 C++ 后端。

模型大小

对于一些研究中的先进模型,您可能需要使用非常大的模型(大小可能达到 GB 级)。由于每个浏览器标签页的内存用量有限,这些模型目前无法在网络浏览器中运行。如需运行这些较大的模型,您可以使用 Node.js 在自己的服务器上高效运行此类模型,并且服务器的硬件达到自己的需求。

IoT

Raspberry Pi 等热门单板计算机均支持 Node.js,这意味着您也可以在此类设备上执行 TensorFlow.js 模型。

速度

Node.js 采用 JavaScript 语言编写,这意味着您可以从即时编译中受益。也就是说,使用 Node.js 时,您可能会经常发现性能有所提升,因为系统会在运行时对之进行优化,特别是对于您可能正在执行的任何预处理操作。此案例研究中就有一个很好的例子,展示了 Hugging Face 如何使用 Node.js 使其自然语言处理模型的性能提升了两倍。

现在,您已了解 TensorFlow.js 的基础知识,它可以运行的环境,以及它具备的一些优势。接下来,我们来运用一些实用功能。

3. 设置系统

在本教程中,我们将使用 Ubuntu。Ubuntu 是一种广受欢迎的 Linux 发行版,许多人都在使用它。如果您选择在基于云的虚拟机上学习本教程,则可以在 Google Cloud Compute Engine 上将 Ubuntu 用作基础映像。

在撰写本文时,我们可以选择 Ubuntu 18.04.4 LTS 的映像来创建新的标准 Compute Engine 实例,而这正是我们将要使用的映像。您当然可以使用自己的机器,甚至可以选择使用其他操作系统,但不同系统之间的安装说明和依赖项可能会有所不同。

安装 TensorFlow(Python 版本)

现在,您可能正尝试转换您找到或将要编写的某个基于 Python 的现有模型,在我们可以从 Python 导出“SavedModel”文件之前,您需要在实例上设置 Python 版 TensorFlow(如果“SavedModel”尚未可供下载)。

通过 SSH 连接到您在上面创建的云端机器,然后在终端窗口中输入以下内容:

终端窗口

sudo apt update
sudo apt-get install python3

这样可确保我们在机器上安装了 Python 3。必须安装 Python 3.4 或更高版本才能使用 TensorFlow。

如需验证是否安装了正确的版本,请输入以下内容:

终端窗口

python3 --version

您应该会看到一些输出内容,其中会显示版本号,例如 Python 3.6.9。如果您看到此值正确无误地打印出来,并且高于 3.4,则可以继续。

接下来,我们将为 Python 3 安装 PIP(Python 的软件包管理器),然后更新它。Type:

终端窗口

sudo apt install python3-pip
pip3 install --upgrade pip

我们还可以通过以下命令验证 pip3 安装:

终端窗口

pip3 --version

在撰写本文时,我们看到执行此命令后,终端中会输出 pip 20.2.3

在安装 TensorFlow 之前,需要将 Python 软件包“setuptools”升级到 41.0.0 或更高版本。运行以下命令,确保其已更新到最新版本:

终端窗口

pip3 install -U setuptools

最后,我们现在可以安装适用于 Python 的 TensorFlow 了:

终端窗口

pip3 install tensorflow

此操作可能需要一段时间才能完成,因此请等待其执行完毕。

我们来检查 TensorFlow 是否已正确安装。在当前目录中创建一个名为 test.py 的 Python 文件:

终端窗口

nano test.py

打开 nano 后,我们可以编写一些 Python 代码来打印已安装的 TensorFlow 版本:

test.py:

import tensorflow as tf
print(tf.__version__)

CTRL + O 将更改写入磁盘,然后按 CTRL + X 退出 nano 编辑器。

现在,我们可以运行此 Python 文件,查看屏幕上打印的 TensorFlow 版本:

终端窗口

python3 test.py

在撰写本文时,我们看到控制台中输出了 2.3.1,这表示我们安装的 TensorFlow Python 版本。

4. 创建 Python 模型

此 Codelab 的下一步将介绍如何创建一个简单的 Python 模型,以展示如何以“SavedModel”格式保存训练后的模型,然后使用 TensorFlow.js 命令行转换器。对于您尝试转换的任何 Python 模型,原理都是类似的,但我们将使此代码保持简单,以便所有人都能理解。

我们来修改在第一部分中创建的 test.py 文件,并将代码更新为如下所示:

test.py:

import tensorflow as tf
print(tf.__version__)

# Import NumPy - package for working with arrays in Python.
import numpy as np

# Import useful keras functions - this is similar to the
# TensorFlow.js Layers API functionality.
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense

# Create a new dense layer with 1 unit, and input shape of [1].
layer0 = Dense(units=1, input_shape=[1])
model = Sequential([layer0])

# Compile the model using stochastic gradient descent as optimiser
# and the mean squared error loss function.
model.compile(optimizer='sgd', loss='mean_absolute_error')

# Provide some training data! Here we are using some fictional data 
# for house square footage and house price (which is simply 1000x the 
# square footage) which our model must learn for itself.
xs = np.array([800.0, 850.0, 900.0, 950.0, 980.0, 1000.0, 1050.0, 1075.0, 1100.0, 1150.0, 1200.0, 1250.0, 1300.0, 1400.0, 1500.0, 1600.0, 1700.0, 1800.0, 1900.0, 2000.0], dtype=float)

ys = np.array([800000.0, 850000.0, 900000.0, 950000.0, 980000.0, 1000000.0, 1050000.0, 1075000.0, 1100000.0, 1150000.0, 1200000.0,  1250000.0, 1300000.0, 1400000.0, 1500000.0, 1600000.0, 1700000.0, 1800000.0, 1900000.0, 2000000.0], dtype=float)

# Train the model for 500 epochs.
model.fit(xs, ys, epochs=500, verbose=0)

# Test the trained model on a test input value
print(model.predict([1200.0]))

# Save the model we just trained to the "SavedModel" format to the
# same directory our test.py file is located.
tf.saved_model.save(model, './')

此代码将训练一个非常简单的线性回归模型,以学习估计所提供的 x(输入)和 y(输出)之间的关系。然后,我们将生成的训练后模型保存到磁盘。请查看内嵌注释,详细了解每一行的作用。

在运行此程序后,如果我们检查目录(通过调用 python3 test.py),现在应该会在当前目录中看到一些新创建的文件和文件夹:

  • test.py
  • saved_model.pb
  • 资产
  • variables

现在,我们已生成 TensorFlow.js 转换器所需的文件,可将此模型转换为在浏览器中运行!

5. 将 SavedModel 转换为 TensorFlow.js 格式

安装 TensorFlow.js 转换器

如需安装转换器,请运行以下命令:

终端窗口

pip3 install tensorflowjs

这很简单。

假设我们使用的是命令行转换器 (tensorflowjs_converter),而不是上面显示的向导版本,我们可以调用以下命令来转换刚刚创建的已保存模型,并明确将参数传递给转换器:

终端窗口

tensorflowjs_converter \
    --input_format=keras_saved_model \
    ./ \
    ./predict_houses_tfjs

这是怎么回事?首先,我们调用刚刚安装的 tensorflowjs_converter 二进制文件,并指定我们尝试转换 Keras 保存的模型。

在上面的示例代码中,您会注意到我们导入了 Keras,并使用其更高级别的层 API 来创建模型。如果您未在 Python 代码中使用 Keras,则可能需要使用其他输入格式:

  • keras - 用于加载 Keras 格式(HDF5 文件类型)
  • tf_saved_model - 用于加载使用 TensorFlow 核心 API(而非 Keras)的模型。
  • tf_frozen_model - 用于加载包含冻结权重的模型。
  • tf_hub - 用于加载从 TensorFlow Hub 生成的模型。

您可以点击此处详细了解其他格式

接下来的 2 个参数用于指定保存的模型所在的文件夹 - 在上面的演示中,我们指定了当前目录,最后指定了希望将转换结果输出到哪个目录,我们在上面将其指定为当前目录中名为“predict_houses_tfjs”的文件夹。

运行上述命令会在当前目录中创建一个名为 predict_houses_tfjs 的新文件夹,其中包含 :

  • model.json
  • Group1-shard1of1.bin

这些文件是我们在 Web 浏览器中运行模型所需的文件。保存这些文件,因为我们将在下一部分中使用它们。

6. 在浏览器中使用转换后的模型

托管转换后的文件

首先,我们必须将生成的 model.json*.bin 文件放在 Web 服务器上,以便我们通过网页访问它们。在此次演示中,我们将使用 Glitch.com,以便您轻松跟进。不过,如果您有 Web 工程背景,也可以选择在当前的 Ubuntu 服务器实例上启动一个简单的 HTTP 服务器来完成此操作。一切由你决定。

将文件上传到 Glitch

  1. 登录 Glitch.com
  2. 使用此链接克隆我们的样板 TensorFlow.js 项目。其中包含一个框架 HTML、CSS 和 JS 文件,这些文件会导入 TensorFlow.js 库,以便我们随时使用。
  3. 点击左侧面板中的“assets”文件夹。
  4. 点击“上传素材资源”,然后选择 group1-shard1of1.bin,即可上传到此文件夹。上传后,该界面现在应如下所示:25a2251c7f165184.png
  5. 如果您点击刚刚上传的 group1-shard1of1.bin 文件,便可以将该网址复制到对应位置。立即复制此路径,如下所示:92ded8d46442c404.png
  6. 现在,在本地机器上使用您偏好的文本编辑器修改 model.json,然后搜索(使用 CTRL+F)其中提及的 group1-shard1of1.bin 文件。

将此文件名替换为您在第 5 步中复制的网址,但删除复制的路径中生成的开头 https://cdn.glitch.com/ 这一错误。

修改后,它应如下所示(请注意,系统已移除开头的服务器路径,因此仅保留了上传后生成的文件名):d5a338f2dc1f31d4.png 7. 现在,点击“assets”(素材资源),然后点击“upload an asset”(上传素材资源)按钮 (重要),将此修改后的 model.json 文件保存并上传到 Glitch。如果您不使用实体按钮,而是采用拖放方式,则文件将作为可修改的文件上传,而不是上传到 CDN。这样一来,文件将不在同一文件夹中,并且当 TensorFlow.js 尝试下载给定模型的二进制文件时,系统会假定相对路径。如果操作正确,您应该会在 assets 文件夹中看到 2 个文件,如下所示:51a6dbd5d3097ffc.png

太棒了!现在,我们可以在浏览器中使用包含实际代码的已保存文件了。

加载模型

现在,我们已经托管了转换后的文件,可以编写一个简单的网页来加载这些文件,并使用它们进行预测。在 Glitch 项目文件夹中打开 script.js,然后将此文件的内容替换为以下内容,前提是您已更改 const MODEL_URL 以指向您在 Glitch 上上传的 model.json 文件所生成的 Glitch.com 链接:

script.js:

// Grab a reference to our status text element on the web page.
// Initially we print out the loaded version of TFJS.
const status = document.getElementById('status');
status.innerText = 'Loaded TensorFlow.js - version: ' + tf.version.tfjs;

// Specify location of our Model.json file we uploaded to the Glitch.com CDN.
const MODEL_URL = YOUR MODEL.JSON URL HERE! CHANGE THIS!';
// Specify a test value we wish to use in our prediction.
// Here we use 950, so we expect the result to be close to 950,000.
const TEST_VALUE = 950.0

// Create an asynchronous function.
async function run() {
    // Load the model from the CDN.
    const model = await tf.loadLayersModel(MODEL_URL);

    // Print out the architecture of the loaded model.
    // This is useful to see that it matches what we built in Python.
    console.log(model.summary());

    // Create a 1 dimensional tensor with our test value.
    const input = tf.tensor1d([TEST_VALUE]);

    // Actually make the prediction.
    const result = model.predict(input);

    // Grab the result of prediction using dataSync method
    // which ensures we do this synchronously.
    status.innerText = 'Input of ' + TEST_VALUE + 
        'sqft predicted as $' + result.dataSync()[0];
}

// Call our function to start the prediction!
run();

MODEL_URL 常量更改为指向 model.json 路径后,运行上述代码会生成如下所示的输出。

c5e8457213058ec3.png

如果我们检查网络浏览器的控制台(按 F12 可在浏览器中打开开发者工具),还可以看到已加载模型的模型说明,其中会显示:

35e79d70dbd66f27.png

将此代码与本 Codelab 开头的 Python 代码进行比较,我们可以确认,这是我们创建的具有 1 个密集输入和 1 个具有 1 个节点的密集层的相同网络。

恭喜!您刚刚在 Web 浏览器中运行了一个转换后的 Python 训练模型!

7. 不转换的模型

有时,编译后使用不太常见操作的更复杂模型将不支持转换。基于浏览器的 TensorFlow.js 版本是对 TensorFlow 的完全重写,因此我们目前不支持 TensorFlow C++ API 具有的所有低层级操作(有数千个)- 不过随着我们的发展以及核心操作变得更加稳定,我们正在逐步添加更多操作。

在撰写本文时,TensorFlow Python 中有一个函数在导出为 SavedModel 时会生成不受支持的 op,该函数是 linalg.diag。如果我们尝试在 Python 中转换使用此功能的已保存模型(Python 支持其生成的相应操作),则会看到类似于以下内容的错误:

5df94fc652393e00.png

在此处,我们可以看到以红色突出显示的部分,即 linalg.diag 调用编译后生成了一个名为 MatrixDiagV3 的操作,但在撰写此 Codelab 时,该操作不受 Web 浏览器中的 TensorFlow.js 支持。

该怎么做?

您可以使用以下两种方法。

  1. 在 TensorFlow.js 中实现此缺失的操作 - 我们是一个开源项目,欢迎大家为新操作等内容做出贡献。请参阅这篇关于为 TensorFlow.js 编写新操作的指南。如果您设法做到了这一点,则可以在命令行转换器中使用 Skip_op_check 标志来忽略此错误,并继续进行转换(它会假定您创建的新版 TensorFlow.js 构建支持缺少的 操作)。
  2. 确定您导出的 savedmodel 文件中的哪个 Python 代码部分产生了不受支持的操作。在少量代码中,这可能很容易找到,但在更复杂的模型中,这可能需要相当多的调查,因为目前没有方法可以在 savedmodel 文件格式中识别生成给定操作的高级 Python 函数调用。不过,找到后,您可以更改为使用其他受支持的方法。

8. 恭喜

恭喜,您已学会初步在网络浏览器中通过 TensorFlow.js 使用 Python 模型!

回顾

在此 Codelab 中,我们学习了如何执行以下操作:

  1. 设置 Linux 环境以安装基于 Python 的 TensorFlow
  2. 导出 Python“SavedModel”
  3. 安装 TensorFlow.js 命令行转换器
  4. 使用 TensorFlow.js 命令行转换器创建所需的客户端文件
  5. 在实际 Web 应用中使用生成的文件
  6. 确定哪些模型不会转化,以及需要实现哪些功能才能让它们在未来实现转化。

后续操作

记得使用 #MadeWithTFJS 标签通知我们,争取在我们的社交媒体呈现甚至参与我们未来的 TensorFlow 活动的机会。我们很想知道您会在浏览器中转化和使用哪些客户端模型!

更多可深入学习的 TensorFlow.js Codelab

参考网站