tensorflow.js GPU崩溃

Question

我的模型获取一个图像并计算一个特定的值，输入层是一个裁剪层，它从图像的顶部和底部移除一些像素。该模型运行良好，但当我更改裁剪层的设置时，例如，从顶部删除25个像素而不是75，浏览器窗口(chrome)会闪烁并输出以下错误：

​

​

注意:就在上面的错误之前，它打印出下面的消息Couldn't parse line number in error，后面是GLSL代码。

如果完全删除裁剪层，也会出现同样的错误。

我使用的是tfjsv3.8.0，但也用v2.0.0进行了测试，结果相似。这是我的模型：

const model = tf.sequential();

// Cropping Layer
model.add(
  tf.layers.cropping2D({
    // If I change 75 to anything below 50, it crashes before completing the first epoch,
    // If this layer is removed, it crashes almost immediately after training starts
    cropping: [
      [75, 25],
      [0, 0]
    ],
    // image height, width, depth
    inputShape: [160, 320, 3]
  })
);

model.add(
  tf.layers.conv2d({
    filters: 16,
    kernelSize: [3, 3],
    strides: [2, 2],
    activation: 'relu',
  })
);

model.add(
  tf.layers.maxPool2d({
    poolSize: [2, 2]
  })
);

model.add(
  tf.layers.conv2d({
    filters: 32,
    kernelSize: [3, 3],
    strides: [2, 2],
    activation: 'relu'
  })
);

model.add(
  tf.layers.maxPool2d({
    poolSize: [2, 2]
  })
);

model.add( tf.layers.flatten());
model.add( tf.layers.dense({ units: 1024, activation: 'relu' }));
model.add( tf.layers.dropout({ rate: 0.25 }));
model.add( tf.layers.dense({ units: 128, activation: 'relu' }));
model.add( tf.layers.dense({ units: 1, activation: 'linear' }));

model.compile({
  optimizer: 'adam',
  loss: 'meanSquaredError',
  metrics: [
    'accuracy',
  ],
});

我是不是做错了什么事？

Answer 1

正如@vladimir-mandic建议的那样，问题最终是GPU耗尽了内存。但是，将WEBGL_DELETE_TEXTURE_THRESHOLD设置为零对我的情况没有帮助。

我花了一段时间来验证这一点，因为它发生在批之间，而且我无法通过tf.memory()在batchEnd上跟踪它，因为在回调过程中，内存已经释放了，否则GPU会在到达这一点之前崩溃。最后我做了两件事来克服这个问题：

• Reduce图像大小：裁剪层有助于不达到“内存不足”状态，从而删除它或减少被裁剪的像素数，从而导致应用程序崩溃。但是由于它也分配了张量，所以我决定通过画布操作重新调整图像的大小，然后将它们输入model.
• Reduce the batchSize:，我一直在使用默认的32 batchSize，直到我将其还原后，我才开始注意到崩溃会消失，这促使我调查了model.fitDataset的内部结构，这就是我如何发现批间内存消耗过大的原因。

正如@vladimir建议的那样，将WEBGL_DELETE_TEXTURE_THRESHOLD设置为0也有助于缓解这个问题，但我没有注意到在我的情况下有什么显著的影响，所以我最终没有使用它。

Answer 2

失去GL上下文意味着GPU内存可能用完了

与实际需要的模型相比，WebGL后端具有很大的开销。

通过强制进行GL内存清理，而不是将其保留在默认情况下，您可能会得到进一步的进展，但这是有代价的(它确实会显著减慢速度)。

tf.ENV.set('WEBGL_DELETE_TEXTURE_THRESHOLD', 0);