基于tensorflow.js设计、训练面向web的神经网络模型的经验

编者按：opencv4nodejs作者Vincent Mühler分享了基于tensorflow.js设计、训练面向web的神经网络模型的经验。

在将现有的一些目标检测、脸部检测、脸部识别的模型移植到tensorflow.js后，我发现有的模型在浏览器中的表现没有达到最优，而另一些模型在浏览器上的表现相当不错。浏览器内的机器学习潜力巨大，tensorflow.js等库给web开发者带来了太多可能性。

然而，直接在浏览器中运行的深度模型也带来了新的挑战和限制，毕竟一些现存的模型不是为在客户端浏览器中运行而设计的，更别说移动端浏览器了。就拿当前最先进的目标检测模型来说：它们通常需要大量计算资源才能以合理的fps运行，更别说实时运行了。另外，对一个简单的web应用来说，分发100MB以上的模型权重至客户端浏览器也不现实。

为web训练高效的深度学习模型

不过不要放弃希望！基于一些基本原则，我们可以创建和训练很不错的模型，为在web环境中运行而优化。信不信由你，实际上我们可以训练相当不错的图像分类模型，甚至是目标检测模型，大小不过几兆，甚至几百K：

这篇文章将给出一些通用的建议，帮助你开始训练自己的卷积神经网络（CNN），以及一些基于tensorflow.js为web应用和移动端浏览器训练CNN的建议。

你和我说要在浏览器里训练深度学习模型？

你也许会好奇：为什么要在浏览器里基于tensorflow.js训练我的模型，而不是直接在自己的机器上基于tensorflow训练模型？你当然可以在自己的机器上训练，特别是如果你的机器配备了NVIDIA显卡。tensorflow.js在底层使用了WebGL加速，所以在浏览器中训练模型的一个好处是可以利用AMD显卡。另外，在浏览器中训练模型，可以更好地保护用户隐私，更容易让用户信任。

网络架构

显然，训练模型之间，首先需要实现模型。通常人们会建议挑选一个久经考验的现有架构，比如YOLO、SSD、ResNet、MobileNet等。

我个人认为，在设计自己的架构时应用现有架构的一些概念是很有价值的，但直接在web应用中引入这些现存架构可能不妥，因为我们想要尺寸较小、推理较快（理想情况下能实时推理）、容易训练的模型。

不管你是打算引入一个现有架构，还是完全从头开始，都可以参考下面的一些建议，在我自己为web设计高效CNN架构时，这些建议帮助很大：

1. 从小型架构开始

牢记，能够取得较好的精确度的网络越小，它进行推理的时间越短，用户下载、缓存模型更容易。此外，较小的模型参数较少，因此在训练时收敛更快。

如果发现当前的网络架构表现不是很好，或者没有达到你想要的精确度，你仍然可以逐渐增加网络的尺寸，例如，增加每层的卷积过滤器数量，或者直接堆叠更多网络层，构造更深的网络。

2. 应用深度可分离卷积

既然我们是要打造一个新模型，毫无疑问我们想要使用深度可分离卷积，而不是传统的2D卷积。深度可分离卷积将常规的卷积操作分为两部分，首先分频道进行卷积，接着应用1x1卷积。和常规卷积操作相比，深度可分离卷积的参数更少，这意味着更少的浮点运算，更容易并行化，推理更快（我曾经碰到过，仅仅将常规卷积替换为深度可分离卷积后，推理就加速了10x），更少的资源占用（在移动设备上意味着性能提升）。此外，因为参数更少，训练所需时间也更短。

MobileNet和Xception应用了深度可分离卷积（tensorflow.js中的MobileNet和PoseNet）。深度可分离卷积是否导致精确度下降大概没有定论，但就我的经验而言，web和移动端模型，毫无疑问应该使用深度可分离卷积。

长话短说：我建议在第一层使用常规的conv2d，毕竟通常而言这一层也没有多少参数，这样可以在提取特征中保留RGB通道之间的关系。

export type ConvParams = {

filter: tf.Tensor4D

bias: tf.Tensor1D

}

exportfunction convLayer(

x: tf.Tensor4D,

params: ConvParams,

stride: [number, number],

padding: string

): tf.Tensor4D {

return tf.tidy(() => {

let out = tf.conv2d(x, params.filter, stride, padding)

out = tf.add(out, params.bias)

return out

})

}

剩下的都用深度可分离卷积，用3 × 3 × channelsin × 1过滤器和1 × 1 × channelsin × channels_out过滤器替换单个核。

export type SeparableConvParams = {

depthwise_filter: tf.Tensor4D

pointwise_filter: tf.Tensor4D

bias: tf.Tensor1D

}

exportfunction depthwiseSeparableConv(

x: tf.Tensor4D,

params: SeparableConvParams,

stride: [number, number],

padding: string

): tf.Tensor4D {

return tf.tidy(() => {

let out = tf.separableConv2d(x, params.depthwise_filter: tf.Tensor4D, params.pointwise_filter, stride, padding)

out = tf.add(out, params.bias)

return out

})

}

所以，原本使用tf.conv2d，形状为[3, 3, 32, 64]的核，现在改用tf.separableConv2d，一个形状为[3, 3, 32, 1]的核加上一个[1, 1, 32, 64]的核。

3. 跳层连接和密集连接块

一旦决定创建较深的网络，很快就会面临训练神经网络最常遇到的问题：梯度消失问题。若干epoch后，损失以极其微小的幅度下降，导致慢得荒谬的训练，甚至是完全收敛不了。

ResNet和DenseNet使用的跳层连接可以缓解创建更深的架构时遇到的梯度消失问题。我们只需将之前的网络层的输出，添加到网络中较深层的输入之中（在应用激活函数之前）：

跳层连接背后的直觉是，梯度不必仅仅通过卷积层（或全连接层）反向传播（这导致梯度渐渐消失）。它们可以添加跳层连接操作，“跳过”一些网络层。

显然，跳层连接要工作，需要输出和输入之间的形状互相匹配。假设我们想要跳层连接网络层A和网络层B，那么A的输出形状需要和B的输入形状匹配。如果你想要创建残差块或密集连接块，只需确保卷积过滤器数目一致，步长为1，补齐方案相同。顺便说下，还有其他方法，比如补齐A的输出，使其满足B的输入的形状，或是连接之前层的特征映射，使得相连的层的深度互相匹配。

刚开始，我摆弄了一下类似ResNet的方法，直接隔层引入跳层连接，如上图所示。不过我很快发现，密集连接块的效果更好，并且应用密集连接块后，模型收敛的时间马上下降了：

这里是一个密集连接块实现的例子，我在face-api.js的68点脸部识别中将其作为基本构件。这个构件使用了4个深度可分离卷积层（注意，第一个密集块的第一个卷积层用了常规卷积），每块的第一个卷积操作步长为2，以缩小输入：

export type DenseBlock4Params = {

conv0: SeparableConvParams | ConvParams

conv1: SeparableConvParams

conv2: SeparableConvParams

conv3: SeparableConvParams

}

exportfunction denseBlock4(

x: tf.Tensor4D,

denseBlockParams: DenseBlock4Params,

isFirstLayer: boolean = false

): tf.Tensor4D {

return tf.tidy(() => {

const out0 = isFirstLayer

? convLayer(x, denseBlockParams.conv0 as ConvParams, [2, 2], 'same')

: depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2], 'same')

as tf.Tensor4D

const in1 = tf.relu(out0) as tf.Tensor4D

const out1 = depthwiseSeparableConv(in1, denseBlockParams.conv1, [1, 1], 'same')

// 第一个连接

const in2 = tf.relu(tf.add(out0, out1)) as tf.Tensor4D

const out2 = depthwiseSeparableConv(in2, denseBlockParams.conv2, [1, 1], 'same')

// 第二个连接

const in3 = tf.relu(tf.add(out0, tf.add(out1, out2))) as tf.Tensor4D

const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv3, [1, 1], 'same')

// 最后一个连接

return tf.relu(tf.add(out0, tf.add(out1, tf.add(out2, out3)))) as tf.Tensor4D

})

}

4. 使用ReLU系列激活函数

除非你有特别的理由，我建议直接使用tf.relu，因为ReLU系列激活函数有助于缓解梯度消失问题。

你也可以试验ReLU的其他变体，例如leaky ReLU，YOLO架构就用了这个：

exportfunction leakyRelu(x: tf.Tensor, epsilon: number) {

return tf.tidy(() => {

const min = tf.mul(x, tf.scalar(epsilon))    

return tf.maximum(x, min)

})

}

或者Mobilenet用的ReLU-6：

exportfunction relu6(x: tf.Tensor) {

return tf.clipByValue(x, 0, 6)

}

训练

有了初始架构之后，就可以开始训练模型了。

5. 如果拿不准主意，直接用Adam

刚开始训练自己的模型的时候，我想知道哪种优化算法是最好的？我从原始的SGD开始，有时候会陷入局部极小值，甚至导致梯度爆炸——模型权重无限增长，逐渐变为NaN.

我并不打算说，对所有问题而言，Adam都是最优选择，但我发现它是训练新模型最简单也最健壮的方式，只需将初始学习率定为0.001，然后以默认参数启动Adam：

const optimizer = tf.train.adam(0.001)

6. 调整学习率

一旦损失没有明显下降，那么我们的模型很可能收敛了（或陷入局部极小值），无法进一步学习了。此时我们也许可以直接停止训练过程，以免模型过拟合（或者尝试不同的架构）。

然而，也有可能还有一点压榨的余地，通过调整（调低）学习率，让模型进一步学习。特别是在整个训练集上的总误差开始震荡（一会儿高一会儿低）的时候，试着降低学习率也许是个好主意。

下图是个例子。从第46个epoch开始，损失值开始震荡，从46个epoch后，将学习率从0.001调低至0.0001，再训练10个epoch，可以进一步降低总误差。

7. 权重初始化

如果你对如何恰当地初始化权重毫无头绪（我刚开始的时候就是这样），那么可以简单地将所有偏置初始化为零（tf.zeros），将权重初始化为从某种正态分布中随机取样的非零值（比如使用tf.randomNormal）。不过我现在更喜欢用glorot正态分布：

const initializer = tf.initializers.glorotNormal()

const depthwise_filter = initializer.apply([3, 3, 32, 1])

const pointwise_filter = initializer.apply([1, 1, 32, 64])

const bias = tf.zeros([64])

8. 打乱输入

训练神经网络的常见建议是在每个epoch前随机化训练样本的出现顺序。我们可以使用tf.utils.shuffle：

exportfunction shuffle(array: any[]|Uint32Array|Int32Array|Float32Array): void

9. 使用FileSaver.js保存模型的快照

FileSaver.js提供了一个saveAs函数，让我们可以储存任意类型的文件至下载文件夹。我们可以用它来保存模型权重：

const weights = newFloat32Array([... model weights, flat array])

saveAs(newBlob([weights]), 'checkpoint_epoch1.weights')

我们也可以存储为json格式，例如保存epoch的累计损失：

const losses = { totalLoss: ... }

saveAs(newBlob([JSON.stringify(losses)]), 'loss_epoch1.json')

调错

在花费大量时间训练模型之前，我们想要确保模型确实能够学习我们预想的东西，并清除任何潜在的错误和bug源。下面的几条建议有助于你避免浪费大量时间训练出一堆垃圾并怀疑人生：

10. 检查输入数据、预处理和后处理逻辑

如果你传入网络的是垃圾，网络返还的也会是垃圾。因此，需要确保输入数据标注正确，网络的输入也符合预期。特别地，如果你实现了随机剪裁、补齐、截成正方形、居中、减去均值之类的预处理逻辑，别忘了可视化预处理之后的输入，以检查输入是否符合预期。另外，我强烈建议单元测试这些步骤。后处理也是一样。

我知道这听起来是一些乏味的额外工作，但它们无疑是有价值的！你不会相信，我花了多少小时尝试搞明白为什么我的目标检测器完全学不会检测面部，直到我逐渐发现，由于错误的剪裁和变形，我的预处理逻辑将输入变成了垃圾。

11. 检查损失函数

在大多数情况下，tensorflow.js都可以提供所需的损失函数。然而，在你需要实现自己的损失函数的情况下，你绝对需要单元测试！不久前我曾经基于tfjs-core API从头构建YOLO v2的损失函数，从头实现损失函数很麻烦，除非你分解问题，并确保每一部分的计算结果符合预期。

12. 先过拟合小数据集

先过拟合训练数据的一个小子集，从而验证损失函数可以收敛，模型确实可以学到有用的东西。一般而言，这是一个好主意。例如，你可以直接从训练数据中选取10到20张图像，并训练若干epoch。损失收敛后，在这10到20张图像上运行推理，并可视化结果：

这是非常重要的一步，有助于消除神经网络实现中各种造成bug的来源，包括预处理逻辑和后处理逻辑在内。因为，如果代码中有不少bug，模型不太可能作出符合期望的预测。

特别是实现自己的损失函数的时候，毫无疑问，你希望在正式开始训练之前确保模型能够收敛。

性能

最后，我想给出一些有助于降低训练时间，以及防止浏览器因内存泄漏而奔溃的建议。

13. 避免明显的内存泄露

除非你在tensorflow.js方面完全是新手，你大概已经知道你需要手动丢弃未使用的张量，以释放它们占用的内存。你可以通过调用tensor.dispose()或将操作封装在tf.tidy块中做到这一点。确保不要因为没有正确地丢弃张量而导致内存泄露，否则你的应用迟早会用尽内存。

识别这类内存泄露很容易。只需在若干次迭代中调用tf.memory()记录日志，以验证张量的数量没有随着每次迭代而异常增长：

14. 调整canvas的尺寸，而不是张量的尺寸

注意，这条建议只适用于当前的tfjs-core（我使用的版本是0.12.14），因为未来的版本可能会修正这一问题。

我知道这也许听起来有点奇怪：为什么不使用tf.resizeBilinear、tf.pad等来重整输入张量的形状，以匹配网络输入要求的形状？这是因为当前tfjs有一个GPU显存泄露的bug（#604）

长话短说，在调用tf.fromPixels将canvas转换成张量前，首先调整canvas的尺寸，以匹配网络输入要求的形状。不这么做的话，你会很快耗尽GPU显存（取决于训练数据中图像尺寸的不同程度）。如果所有训练图像尺寸一致，不会有什么问题，但如果尺寸不一致，你可以使用以下代码调整尺寸：

exportfunction imageToSquare(img: HTMLImageElement | HTMLCanvasElement, inputSize: number): HTMLCanvasElement {

const dims = img instanceof HTMLImageElement

? { width: img.naturalWidth, height: img.naturalHeight }

: img

const scale = inputSize / Math.max(dims.height, dims.width)

const width = scale * dims.width

const height = scale * dims.height

const targetCanvas = document.createElement('canvas')

targetCanvas .width = inputSize

targetCanvas .height = inputSize

targetCanvas.getContext('2d').drawImage(img, 0, 0, width, height)

return targetCanvas

}

15. 找到最优的batch大小

不要把batch大小设得过大！尝试不同的batch大小，测量反向传播所需的时间。最优batch大小取决于GPU、输入尺寸、网络复杂度。甚至某些情况下，不使用batch，逐一输入才是最好的。

如果拿不准，我会将batch大小设为1（也就是不使用batch）。我个人发现在一些情形下，增加batch大小对性能提升没什么帮助，不过，在另一些情形下，我发现，在相当小的网络上，为112 × 112的图像创建16到24的batch能带来大约1.5-2.0的速度提升。

16. 缓存、离线存储、Indexeddb

训练图像（和标签）也许很多，取决于图像的尺寸和数量，可能达到1GB甚至更多。由于在浏览器下无法直接从磁盘读取图像，我们需要使用一个文件代理，比如一个简单的express服务器，托管训练数据，然后让浏览器通过文件代理获取数据。

显然，这很不高效，不过，别忘了，在浏览器中训练时，如果数据集足够小，我们大概会把所有数据放在内存中，这显然也不怎么高效。刚开始，我尝试增加浏览器的缓存大小，以直接将所有数据缓存到磁盘上，不过看起来新版的Chrome和FireFox都不再支持这个功能了。

我最终决定使用Indexddb，这是一个浏览器中的数据库，可以用来存储整个训练集和测试集。上手Indexeddb很容易，只需几行代码，就可以以键值存储的格式储存和查询所有数据。在Indexeddb中，我们可以将标签存储为json对象，将图像数据存储为blob。

查询Indexeddb相当快，至少比反复从文件代理查询并获取文件要快。另外，将数据转移到Indexeddb之后，整个训练过程完全可以离线进行，也就是说，之后的训练过程不再需要文件代理服务器了。

17. 异步损失报告

这个一个非常有效的简单技巧，有助于大量减少训练的迭代次数。如果我们想要获取optimizer.minimize返回的损失张量的值，以了解训练过程中的损失值，我们希望避免等待每次迭代的CPU和GPU同步数据。相反，我们希望迭代异步地报告损失：

const loss = optimizer.minimize(() => {

const out = net.predict(someInput)

const loss = tf.losses.meanSquaredError(

groundTruth,

out,

tf.Reduction.MEAN

)

return loss

}, true)

loss.data().then(data => {

const lossValue = data[0]

window.lossValues[epoch] += (window.lossValues[epoch] || 0) + lossValue

loss.dispose()

})

别忘了，现在损失是异步报告的，所以如果我们想要把每个epoch末的总体损失保存到一个文件，我们需要等待最后一个promise得到满足。我通常直接使用setTimeout在每个epoch完成后的10秒之后再记录总体损失。

if (epoch !== startEpoch) {

// 丑陋的等待最后一个promise满足的方法

const previousEpoch = epoch - 1

setTimeout(() => storeLoss(previousEpoch, window.losses[previousEpoch]), 10000)

}

成功训练模型之后

18. 权重量化

完成模型训练并对模型的表现满意后，我建议应用权重量化以压缩模型大小。通过量化模型权重，可以将模型大小压缩至1/4！尽可能地压缩模型大小是很关键的，因为它有助于快速地传输模型权重至客户端应用，特别是这样的压缩几乎没有什么代价。

所以，别忘了参考我写的tensorflow.js的权重量化指南：https://itnext.io/shrink-your-tensorflow-js-web-model-size-with-weight-quantization-6ddb4fcb6d0d