TensorFlow.js 在浏览器和 Node.js 中运行,并且在这两个平台上都有许多不同的可用配置。每个平台都有一组独特的注意事项,会影响应用程序的开发方式。
在浏览器中,TensorFlow.js 支持移动设备和桌面设备。每个设备都有一组特定的约束,例如可用的 WebGL API,这些约束会自动为您确定和配置。
在 Node.js 中,TensorFlow.js 支持直接绑定到 TensorFlow API 或使用较慢的普通 CPU 实现运行。
环境
执行 TensorFlow.js 程序时,特定配置称为环境。环境由单个全局后端以及控制 TensorFlow.js 微粒度功能的一组标志组成。
后端
TensorFlow.js 支持多种不同的后端,这些后端实现张量存储和数学运算。在任何给定时间,只有一个后端处于活动状态。大多数情况下,TensorFlow.js 会根据当前环境自动为您选择最佳后端。但是,有时了解正在使用的后端以及如何切换后端非常重要。
要查找您正在使用的后端
console.log(tf.getBackend());
如果您想手动更改后端
tf.setBackend('cpu');
console.log(tf.getBackend());
WebGL 后端
WebGL 后端“webgl”目前是浏览器中最强大的后端。该后端比普通 CPU 后端快 100 倍。张量存储为 WebGL 纹理,数学运算在 WebGL 着色器中实现。以下是一些在使用此后端时需要了解的有用信息:\
避免阻塞 UI 线程
调用操作时,例如 tf.matMul(a, b),将同步返回结果 tf.Tensor,但矩阵乘法的计算可能尚未真正准备好。这意味着返回的 tf.Tensor 只是一个计算句柄。当您调用 x.data()
或 x.array()
时,值将在计算实际完成时解析。这使得使用异步 x.data()
和 x.array()
方法比其同步对应方法 x.dataSync()
和 x.arraySync()
更加重要,以避免在计算完成时阻塞 UI 线程。
内存管理
使用 WebGL 后端时的一个注意事项是需要显式内存管理。WebGLTextures(张量数据最终存储的位置)不会被浏览器自动垃圾回收。
要销毁 tf.Tensor
的内存,可以使用 dispose()
方法
const a = tf.tensor([[1, 2], [3, 4]]);
a.dispose();
在应用程序中将多个操作链接在一起非常常见。保存对所有中间变量的引用以将其释放会降低代码可读性。为了解决这个问题,TensorFlow.js 提供了一个 tf.tidy()
方法,该方法会在执行函数后清理所有未返回的 tf.Tensor
,类似于函数执行时局部变量的清理方式
const a = tf.tensor([[1, 2], [3, 4]]);
const y = tf.tidy(() => {
const result = a.square().log().neg();
return result;
});
精度
在移动设备上,WebGL 可能只支持 16 位浮点纹理。但是,大多数机器学习模型都是使用 32 位浮点权重和激活进行训练的。当为移动设备移植模型时,这会导致精度问题,因为 16 位浮点数只能表示范围为 [0.000000059605, 65504]
的数字。这意味着您应该注意模型中的权重和激活不要超过此范围。要检查设备是否支持 32 位纹理,请检查 tf.ENV.getBool('WEBGL_RENDER_FLOAT32_CAPABLE')
的值,如果为 false,则该设备只支持 16 位浮点纹理。您可以使用 tf.ENV.getBool('WEBGL_RENDER_FLOAT32_ENABLED')
检查 TensorFlow.js 当前是否正在使用 32 位纹理。
着色器编译和纹理上传
TensorFlow.js 通过运行 WebGL 着色器程序在 GPU 上执行操作。这些着色器在用户请求执行操作时被延迟地组装和编译。着色器的编译发生在 CPU 上的主线程上,可能很慢。TensorFlow.js 会自动缓存已编译的着色器,使对具有相同形状的输入和输出张量的相同操作的第二次调用快得多。通常,TensorFlow.js 应用程序会在应用程序的生命周期中多次使用相同的操作,因此机器学习模型的第二次传递速度要快得多。
TensorFlow.js 还将 tf.Tensor 数据存储为 WebGLTextures。当一个 tf.Tensor
被创建时,我们不会立即将数据上传到 GPU,而是将数据保留在 CPU 上,直到 tf.Tensor
在操作中使用。如果 tf.Tensor
第二次使用,数据已经存在于 GPU 上,因此没有上传成本。在典型的机器学习模型中,这意味着权重在模型的第一次预测期间上传,而模型的第二次传递速度将快得多。
如果您关心模型或 TensorFlow.js 代码的第一次预测的性能,我们建议在使用真实数据之前,通过传递相同形状的输入张量来预热模型。
例如
const model = await tf.loadLayersModel(modelUrl);
// Warmup the model before using real data.
const warmupResult = model.predict(tf.zeros(inputShape));
warmupResult.dataSync();
warmupResult.dispose();
// The second predict() will be much faster
const result = model.predict(userData);
Node.js TensorFlow 后端
在 TensorFlow Node.js 后端“node”中,TensorFlow C API 用于加速操作。如果可用,这将使用机器的可用硬件加速,例如 CUDA。
在这个后端,就像 WebGL 后端一样,操作同步返回 tf.Tensor
。但是,与 WebGL 后端不同,操作在您获得张量之前完成。这意味着对 tf.matMul(a, b)
的调用将阻塞 UI 线程。
因此,如果您打算在生产应用程序中使用它,您应该在工作线程中运行 TensorFlow.js 以避免阻塞主线程。
有关 Node.js 的更多信息,请参阅本指南。
WASM 后端
TensorFlow.js 提供了一个 WebAssembly 后端 (wasm
),它提供 CPU 加速,可以用作普通 JavaScript CPU (cpu
) 和 WebGL 加速 (webgl
) 后端的替代方案。要使用它
// Set the backend to WASM and wait for the module to be ready.
tf.setBackend('wasm');
tf.ready().then(() => {...});
如果您的服务器在不同的路径或不同的名称上提供 .wasm
文件,请在初始化后端之前使用 setWasmPath
。有关更多信息,请参阅 README 中的 "使用捆绑器" 部分
import {setWasmPath} from '@tensorflow/tfjs-backend-wasm';
setWasmPath(yourCustomPath);
tf.setBackend('wasm');
tf.ready().then(() => {...});
为什么选择 WASM?
WASM 于 2015 年推出,作为一种新的基于 Web 的二进制格式,为用 JavaScript、C、C++ 等编写的程序提供了一个编译目标,以便在 Web 上运行。WASM 已被 支持 Chrome、Safari、Firefox 和 Edge 自 2017 年以来,并得到 全球 90% 的设备 的支持。
性能
WASM 后端利用 XNNPACK 库 来优化神经网络运算符的实现。
与 JavaScript 相比:对于浏览器来说,WASM 二进制文件通常比 JavaScript 捆绑包加载、解析和执行的速度快得多。JavaScript 是动态类型的,并且进行垃圾回收,这会导致运行时速度变慢。
与 WebGL 相比:WebGL 比 WASM 对于大多数模型来说更快,但对于小型模型,WASM 由于执行 WebGL 着色器的固定开销,可以胜过 WebGL。“何时应该使用 WASM”部分将在下面讨论做出此决定的启发式方法。
可移植性和稳定性
WASM 具有可移植的 32 位浮点算术,在所有设备上提供精度一致性。另一方面,WebGL 是特定于硬件的,不同的设备可能具有不同的精度(例如,在 iOS 设备上回退到 16 位浮点数)。
与 WebGL 一样,WASM 也得到所有主要浏览器的官方支持。与 WebGL 不同,WASM 可以在 Node.js 中运行,并且可以在服务器端使用,而无需编译任何本机库。
何时应该使用 WASM?
模型大小和计算需求
一般来说,当模型较小或您关心缺乏 WebGL 支持 (OES_texture_float
扩展) 或 GPU 性能较低的低端设备时,WASM 是一个不错的选择。下图显示了 2018 年 MacBook Pro 上 Chrome 中 5 个我们官方支持的 模型 在 WebGL、WASM 和 CPU 后端上的推理时间(截至 TensorFlow.js 1.5.2)
较小的模型
模型 | WebGL | WASM | CPU | 内存 |
---|---|---|---|---|
BlazeFace | 22.5 毫秒 | 15.6 毫秒 | 315.2 毫秒 | .4 MB |
FaceMesh | 19.3 毫秒 | 19.2 毫秒 | 335 毫秒 | 2.8 MB |
较大的模型
模型 | WebGL | WASM | CPU | 内存 |
---|---|---|---|---|
PoseNet | 42.5 毫秒 | 173.9 毫秒 | 1514.7 毫秒 | 4.5 MB |
BodyPix | 77 毫秒 | 188.4 毫秒 | 2683 毫秒 | 4.6 MB |
MobileNet v2 | 37 毫秒 | 94 毫秒 | 923.6 毫秒 | 13 MB |
上表显示,WASM 在所有模型中比普通 JS CPU 后端快 10-30 倍,并且在较小的模型(如 BlazeFace)中与 WebGL 相当,该模型很轻量级(400KB),但具有相当数量的操作(~140)。鉴于 WebGL 程序每个操作执行都有固定的开销,这解释了为什么像 BlazeFace 这样的模型在 WASM 上更快。
这些结果会因您的设备而异。确定 WASM 是否适合您的应用程序的最佳方法是在我们的不同后端上对其进行测试。
推理与训练
为了解决预训练模型部署的主要用例,WASM 后端开发将优先考虑推理而不是训练支持。请参阅 最新支持的操作列表,以及 告知我们您的模型是否具有不支持的操作。对于训练模型,我们建议使用 Node(TensorFlow C++)后端或 WebGL 后端。
CPU 后端
CPU 后端“cpu”是性能最低的后端,但它是最简单的。所有操作都在普通 JavaScript 中实现,这使得它们不太可并行化。它们还会阻塞 UI 线程。
这个后端对于测试或 WebGL 不可用的设备非常有用。
标志
TensorFlow.js 有一组环境标志,这些标志会自动评估并确定当前平台上的最佳配置。这些标志大多是内部的,但一些全局标志可以通过公共 API 控制。
tf.enableProdMode():
启用生产模式,这将删除模型验证、NaN 检查和其他正确性检查,以换取性能。tf.enableDebugMode()
: 启用调试模式,这将记录执行的每个操作到控制台,以及运行时性能信息,例如内存占用和总内核执行时间。请注意,这会极大地减慢您的应用程序速度,请不要在生产环境中使用它。