用浏览器训练Tensorflow.js模型的18个技巧

在移植现有模型（除tensorflow.js）进行物体检测、人脸检测、人脸识别后，我发现一些模型不能以最佳性能发挥。而tensorflow.js在浏览器中表现相当不错，如果你想见证浏览器内部机器学习的潜力以及tensorflow.js为我们的Web开发人员提供的所有可能性，我个人建议你可以尝试下。

但是，由于深度学习模型无法直接在浏览器中运行，因为这些模型不是专为在浏览器中运行而设计的，更不用说在移动端了。以现有技的物体探测器为例：它们通常需要大量的计算资源才能以合理的fps运行，更不用说以实时速度运行了。此外，在简单的Web应用程序中将100MB+模型权重加载到客户端浏览器是根本不可行的。

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

但让我告诉你，我们能够构建和训练相当不错的模型，这些模型通过设计一些基本原则可以在Web环境中运行进行。信不信由你：我们可以训练相当不错的图像分类-甚至物体检测模型，最终只有几兆字节大小甚至只有几千字节：

在本文中，我想给你一些关于开始训练你自己的卷积神经网络（CNN）的一般技巧，还有一些技巧是直接针对在浏览器中使用张量流训练CNN网络。

现在你可能想知道：为什么我可以在浏览器中使用tensorflow.js训练我的模型，在我的机器上如何使用tensorflow训练它们？当然，如果你的机器配备了NVIDIA卡，你也可以这样做。在浏览器中训练深度学习的一个巨大的优势，它的引擎盖下利用WebGL的，这就意味着你不需要NVIDIA GPU训练的模式，而是AMD GPU上训练深度学习模型。

因此，如果你的机器配备了NVIDIA卡，你可以简单地采用标准张量流方法（在这种情况下，你可以在python中编写训练代码）并忽略浏览器的提示。现在，让我们开始吧！

网络架构

在开始训练我们自己的图像分类器、对象检测器之前，我们必须首先实现网络架构。我通常建议选择现有的架构，例如Yolo、SSD、ResNet、MobileNet等。

就个人而言，我认为在你自己的架构中使用这些架构所采用的一些概念是很有价值的。然而，正如我最初指出的那样，简单地采用这些架构不会让我们的模型是体积小，推理快和易于训练成为可能。

无论你是想要适应现有架构还是从头开始，我都想给你以下建议，这有助于为Web设计高效的CNN架构：

1.从小型网络架构开始！

记住，我们的网络越小，同时仍能在解决问题时获得良好的准确性，它在推理时间内执行的速度就越快，客户端下载和缓存该模型就越容易。此外，较小的模型具有较少的参数，因此在训练时会更快地收敛。

如果你发现当前的网络架构性能不佳，或者达不到准确性水平，你仍然希望是它，你可以逐步增加网络的大小，例如通过增加每层卷积滤波器的数量或堆叠更多层简单地使你的网络更深入。

2.采用深度可分的卷积！

由于我们正在训练一个新模型，我们希望明确使用深度可分卷积而不是普通2D卷积。深度可分离卷积将常规卷积运算分成深度卷积，然后是逐点（1x1）卷积。与常规卷积操作相比，它们具有更少的参数，这将使用更少的浮点运算并且更容易并行化，这意味着推断将更快（我通过简单地替换常规卷积来进行推断的速度提升高达10倍）和较少的资源消耗。此外，因为它们具有较少的参数，所以训练它们所花费的时间较少。

在MobileNet和Xception中采用了深度可分离卷积的思想，你可以在MobileNet和PoseNet的tensorflow.js模型中找到它们。深度可分离卷积是否导致模型不太准确可能是一个公开的辩论，但根据我的经验，它们绝对是网络模型的方式。

长话短说：我建议在你的第一层使用常规的没有那么多的参数的conv2d操作，以保留提取的特征中的RGB通道之间的关系。

export type ConvParams = {
filter: tf.Tensor4D
bias: tf.Tensor1D
}
export function 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
})
}

对于其余的卷积，只需使用深度可分离的卷积。因此，我们将使用3x3xchannels_inx1深度滤波器和1x1x channels_in x channels_out逐点滤波器，而不是单个内核。

export type SeparableConvParams = {
depthwise_filter: tf.Tensor4D
pointwise_filter: tf.Tensor4D
bias: tf.Tensor1D
}
export function 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.跳过连接（Skip connections）和密集连接（Densely Connected）的块

一旦决定建立更深层的网络，很快就面临着训练神经网络最常见的问题之一：梯度消失问题。在一些时期之后，损失只会在非常微小的步骤中减少，这会增加训练时间或者导致模型不收敛。

ResNet和DenseNet中使用的跳过连接允许它们构建更深层的体系结构，同时减轻梯度消失问题。我们所要做的就是在应用激活函数之前，将先前层的输出添加到位于网络中更深层的层输入中：

跳过连接

跳过连接工作，通过快捷方式连接图层。这种技术背后的本质是，梯度不必仅通过卷积（或完全连接）反向传播，这导致梯度一旦到达网络的早期层就会减少。它们可以通过跳过连接的添加操作来“skip”层。

显然，假设你想要将A层与B层连接，A的输出形状必须与B的输入形状相匹配。如果你想构建残差或密集连接的块，只需确保在该块的卷积中保持相同数量的滤波器，并使用相同的填充保持1的步幅。正如旁注一样，也有不同的方法，它们填充A的输出，使其与输入B的形状匹配，或者连接来自先前层的特征映射，使得连接层的深度再次匹配。

起初，我尝试使用类似ResNet的方法，只需在其他层之间引入跳过连接，如上图所示，但很快就发现，密集连接的块工作得更好，并且可以很快的收敛：

密集块

这是一个密集块实现的例子，我用它作为面部标志检测器的基本构建块face-api.js。其中一个块涉及4个可深度分离的卷积层（注意，第一个密集块的第一个卷积是常规卷积），每个块的第一个卷积运算使用2的步幅来缩小输入：

export type DenseBlock4Params = {
conv0: SeparableConvParams | ConvParams
conv1: SeparableConvParams
conv2: SeparableConvParams
conv3: SeparableConvParams
}


export function 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')


// first join
const in2 = tf.relu(tf.add(out0, out1)) as tf.Tensor4D
const out2 = depthwiseSeparableConv(in2, denseBlockParams.conv2, [1, 1], 'same')


// second join
const in3 = tf.relu(tf.add(out0, tf.add(out1, out2))) as tf.Tensor4D
const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv3, [1, 1], 'same')


// final join
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架构中使用：

export function leakyRelu(x: tf.Tensor, epsilon: number) {
return tf.tidy(() => {
const min = tf.mul(x, tf.scalar(epsilon)) 
return tf.maximum(x, min)
})
}

或者正在移动网络使用的ReLU-6：

export function relu6(x: tf.Tensor) {
return tf.clipByValue(x, 0, 6)
}

训练

一旦我们完成初始架构，我们就可以开始训练我们的模型了。

5.如果有疑问，只需使用Adam Optimizer！

当第一次开始训练自己的模型时，我想知道哪种优化器最好？我一开始使用普通的SGD，它似乎有时会陷入局部最小值中，甚至导致梯度爆炸，以至于模型权重无限增长，最终导致NaNs。

我并不是说，Adam是所有问题的最佳选择，但我发现它是训练新模型最简单且最强大的方法，只需使用默认参数和学习率为0.001的Adam开始：

const optimizer = tf.train.adam(0.001)

6.调整学习率

一旦损失没有显着下降，很可能，我们的模型确实收敛，并且无法进一步学习。此时我们不妨停止训练过程，以防止我们的模型不会出现过度拟合。

但是，你可以通过调整（降低）学习率来避免它发生。特别是如果在训练集上计算的总损失开始振荡，这表明尝试降低学习率可能是个好主意。

下面是一个示例，显示了训练面部标志模型时整体误差的图表。在46epoch，损失值开始振荡。正如你所看到的那样，继续训练从46epoch的检查点再学习10个以上，学习率为0.0001而不是0.001，这能够进一步降低整体误差：

7.权重初始化

如果你对如何正确初始化模型权重没有任何线索：一个简单的经验法则，从某种正态分布中得出，用零初始化所有偏差（tf.zeros（shape））和你的权重（卷积的核和全连接层的权重）与非零值。例如，你可以简单地使用tf.randomNormal（shape），但是现在我更喜欢使用glorot正态分布，这在tfjs-layers中是可用的，如下所示：

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.随机你的输入！`

训练神经网络的一个常见建议是通过在每个时期开始时对输入进行混洗来随机化训练样本。我们可以使用tf.utils.shuffle来实现这个目的：

/** Shuffles the array using Fisher-Yates algorithm. */
export function shuffle(array: any[]|Uint32Array|Int32Array|Float32Array): void

`9.使用FileSaver.js保存模型检查点`

由于我们在浏览器中训练我们的模型，你现在可能会问自己：我们如何在训练时自动保存模型权重的检查点？我们可以使用FileSaver.js，该脚本公开了一个名为saveAs的函数，我们可以使用它来存储任意类型的文件，这些文件最终会出现在我们的下载文件夹中。

这样我们就可以保存模型权重：

const weights = new Float32Array([... model weights, flat array])
saveAs(new Blob([weights]), 'checkpoint_epoch1.weights')

甚至是json文件：

const losses = { totalLoss: ... }
saveAs(new Blob([JSON.stringify(losses)]), 'loss_epoch1.json')

`排除故障`

在花费大量时间训练模型之前，你需要确保你的模型实际上要学习是什么，并消除任何潜在的错误来源。如果你不考虑以下提示，你可能会浪费你的时间在训练垃圾上：

`10.检查输入数据，预处理和后处理逻辑！`

如果你将垃圾传递到你的网络，它定会把垃圾扔回你身边。因此，请确保你的输入数据标记正确，并确保你的网络输入符合你的预期。特别是如果你已经规定了一些预处理逻辑，如随机裁剪、填充、平方、居中、平均减法或其他什么，请确保在预处理后进行可视化输入。此外，我强烈建议单元测试这些步骤。

这听起来像是一项繁琐的额外工作，但它很重要！

`11.检查你的损失函数！`

现在在大多数情况下，tensorflow.js为你提供了你需要的损失函数。但是，如果你需要实现自己的损失函数，你绝对需要进行单元测试！不久前，我从头开始使用tfjs-core API实现了Yolo v2 dropout函数，以便为网络训练yolo对象检测器。

`12.先装上一个小型数据集！`

通常，最好是在训练数据的一小部分上过度拟合，以验证损失是否正在收敛，以及你的模型实际上是否在学习一些有用的东西。因此，你应该只选择10到20张训练数据的图像并训练一些时期。一旦损失收敛，对这10到20张图像进行推断并可视化结果：

这是一个非常重要的步骤，它将帮助你消除网络实施中的各种错误来源、前后处理逻辑。

特别是，如果你正在实现自己的损失函数，你肯定要确保，你的模型能够在开始训练之前收敛！

`性能`

最后，我想给你一些建议，通过考虑一些基本原则，这将有助于你尽可能地减少训练时间并防止浏览器因内存泄漏而崩溃。

`13.防止明显的内存泄漏`

除非你是tensorflow.js的新手，否则你可能已经知道，我们必须手动处理未使用的张量来释放内存，方法是调用tensor.dispose（）或将我们的操作包装在tf.tidy块中。确保由于未正确处理张量而导致没有此类内存泄漏，否则你的应用程序迟早会耗尽内存。

识别这些类型的内存泄漏非常简单，只需记录tf.memory（）几次迭代即可验证，每次迭代时张量的数量不会无意中增长：

`14.调整Canvases大小而不是你的张量！`

注意，以下语句仅在tfjs-core的当前状态时有效，直到最终得到修复。

这可能听起来有点奇怪：为什么不使用tf.resizeBilinear、tf.pad等将输入张量重塑为所需的网络输入形状？tfjs目前有一个未解决的问题，说明了这个问题。

TLDR：在调用tf.fromPixels之前，要将Canvaes转换为张量，请调整Canvaes的大小，使其具有网络接受的大小，否则你将快速耗尽GPU内存，具体取决于各种不同的输入大小。如果你的训练图像大小都相同，那么这个问题就不那么严重了，但是如果你必须明确调整它们的大小，你可以使用下面的代码片段：

export function 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.确定最佳批量大小`

不要过分批量输入！尝试不同的批量大小并测量反向传播所需的时间。最佳批量大小显然取决于你的GPU统计信息，输入大小以及网络的复杂程度。在某些情况下，你根本不想批量输入。

如果有疑问的话，我会一直使用1的批量大小。我个人认为，在某些情况下，增加批量大小对性能没有任何帮助，但在其他情况下，我可以看到整体加速的因素通过创建大小为16到24的批次，在相当小的网络尺寸下输入图像大小为112x112像素，大约1.5-2.0左右。

`16.缓存、离线存储、Indexeddb`

我们的训练图像可能相当大，可能高达1GB甚至更大，具体取决于图像的大小和数量。由于我们不能简单地在浏览器中从磁盘读取图像，我们将使用文件代理（可能是一个简单的快速服务器）来托管我们的训练数据，浏览器将获取每个数据项。

显然，这是非常低效的，但是在浏览器中进行训练时我们必须记住这一点，如果你的数据集足够小，你可能会尝试将整个数据保存在内存中，但这显然也不是很有效。最初，我试图增加浏览器缓存大小以简单地将整个数据缓存在磁盘上，但这在以后的Chrome版本中似乎不再起作用，而且我也没有运气FireFox。

最后，我决定只使用Indexeddb，这是一个浏览器数据库，你可能不太熟悉，我们可以利用它来存储我们的整个训练和测试数据集。Indexeddb入门非常简单，因为我们基本上只需几行代码即可将整个数据存储和查询为键值存储。使用Indexeddb，我们可以方便地将标签存储为普通的json对象，将我们的图像数据存储为blob。看看这篇博文，很好地解释了如何在Indexeddb中保存图像数据和其他文件。

查询Indexeddb是非常快的，至少我发现查询每个数据项的速度要快一些，而不是一遍又一遍地从代理服务器中获取文件。此外，在将数据移动到Indexeddb之后，技术上的训练现在完全脱机，这意味着我们可能不再需要代理服务器了。

`17.异步丢失报告`

这是一个简单但非常有效的提示，它帮助我减少了训练时的迭代次数。主要的作用是，如果我们想要检索由optimizer.minimize返回的损失张量的值，我们肯定会这样做，因为我们想要在训练时跟踪我们的损失，我们希望避免等待损失返回的lose.data（）及防止等待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的末尾将整体损失保存到文件中，我们将不得不等待最后的解决方案。我通常只是通过使用setTimeout在一个epoch完成后10秒左右保存整体损失值来解决这个问题：

if (epoch !== startEpoch) {
// ugly hack to wait for loss datas for that epoch to be resolved
const previousEpoch = epoch - 1
setTimeout(() => storeLoss(previousEpoch, window.losses[previousEpoch]), 10000)
}