tensorflow.js GPU 崩溃

Question

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

注意：在出现上述错误之前，它会打印出以下消息 Couldn't parse line number in error，后跟似乎是 GLSL 代码的内容。

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

我使用的是 tfjs v3.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

丢失 GL 上下文意味着您可能运行 GPU 内存不足
与模型实际需要的相比，WebGL 后端具有巨大的开销

通过强制 GL 内存清理而不是将其保留为默认值，您可能会更进一步，但这是有代价的（它确实会显着减慢速度）

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

Answer 2

正如@vladimir-mandic 所建议的那样，最终问题确实是 GPU 运行内存不足。但是将 WEBGL_DELETE_TEXTURE_THRESHOLD 设置为零对我的情况没有帮助。

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

减小图像尺寸： 裁剪层有助于不达到 'out of memory' 状态，因此删除它或减少裁剪的像素数会导致应用程序坠毁。但由于它也分配张量，我决定在将它们输入模型之前通过 canvas 操作调整图像大小。
减少 batchSize： 我一直在使用默认的 batchSize 32，直到我减少它，我才开始注意到崩溃会消失，这导致我调查了 model.fitDataset 的内部结构，这就是我发现批次之间内存消耗过多的原因。

正如@vladimir-mandic 推荐的那样，将 WEBGL_DELETE_TEXTURE_THRESHOLD 设置为 0 应该也有助于缓解这个问题，但我没有注意到对我的情况有任何显着影响，所以我最终没有使用它。

tensorflow.js GPU 崩溃

tensorflow.js GPU crashing

tensorflow

machine-learning

tensorflow.js