tensorflow.js GPU crashing - tensorflow

Question

My model takes an image and calculates a certain value, the input layer is a cropping layer that removes a number of pixels from the top and bottom of the image. The model works quite well, but when I change the settings of the cropping layer, say, remove 25 pixels from the top instead of 75, the browser window (chrome) flickers and outputs the following error:
NOTE: Right before the above error, it prints out the following message Couldn't parse line number in error followed by what appears to be GLSL code.
The same error shows up if I remove the cropping layer altogether.
I'm using tfjs v3.8.0 but also tested with v2.0.0 with a similar outcome. This is my model:
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',
],
});
Am I doing something obviously wrong?

Answer 1

Lost GL context means you've likely run out of GPU memory
WebGL backend has massive overhead compared to what model actually needs
You might get a bit further with forcing GL memory cleanup instead of leaving it at defaults, but it comes at a cost (it does slow things down significantly)
tf.ENV.set('WEBGL_DELETE_TEXTURE_THRESHOLD', 0);

Answer 2

The problem did end up being the GPU running out of memory as #vladimir-mandic suggested. But setting WEBGL_DELETE_TEXTURE_THRESHOLD to zero did not help in my case.
It took me a while to verify this because it happened between batches and I could not track it via tf.memory() at batchEnd because in the callback the memory had already been released OR the GPU would crash before reaching that point. I ended up doing two things to overcome this problem:
Reduce the image size: The cropping layer was helping not to reach the 'out of memory' state, hence removing it or reducing the number of pixels cropped was causing the application to crash. But since it also allocates tensors, I decided to resize the image via canvas manipulations BEFORE feeding them into the model.
Reduce the batchSize: I had been working with the default batchSize of 32 and it was only until I reduced it that I started to notice the crash would go away, this led me to investigate the internals of model.fitDataset and that's how I found the excessive memory consumption between batches.
As #vladimir-mandic recommended, setting WEBGL_DELETE_TEXTURE_THRESHOLD to 0 SHOULD also help alleviate this issue, but I didn't notice any significant effect in my case, so I didn't end up using it.