Warm up layers models when they are loaded.

[nsthorat]

If we pass a tensor of zeros to the model it will compile all the shader programs and upload weights to the GPU, thus making the first inference much faster.

This should be an easy win.

[bileschi]

Can you explain a little more? Something like calling model.predict internally as the last step in model.compile ?

[nsthorat]

Actually I was thinking model.predict(tf.zeros(inputLayer.shape)); as the last step of tf.loadModel.

[bileschi]

Seems reasonable. Quick question: In general a model might have input layer(s) with incomplete shapes, say [?, 100, 100], typically for batch. If we guess the batch size, and are incorrect, does the warm-up behavior still work, or does the GPU memory need to be re-allocated because of the shape change?

[dsmilkov]

Good q. In that case, we should warmup with the most common case in inference, batchSize=1. If we happen to be wrong, some of the GPU programs will be re-compiled, but that's fine.

[nsthorat]

We should warm up with batchSize = 1 since it will be faster. Having batchSize > 1 likely will not change which programs are compiled / weights are uploaded. It's also likely that users will be using a batchSize of 1 for inference, anyways.

[rikkitook]

Little side note: we noticed that on the full tensorflow with gpu enabled, consecutive calls with different batch sizes leads to session reheating. Like, for example, first call with batchsize = 1 is long, first call with batchsize = 5 is also long on the same session

[gaikwadrahul8]

Hi, @nsthorat

Thank you for opening this issue for tracking purposes. Since this issue has been open for a long time, the code/debug information for this issue may not be relevant with the current state of the code base.

The TFJs team is constantly improving the framework by fixing bugs and adding new features. We suggest you try the latest TFJs version with the latest compatible hardware configuration which could potentially resolve the issue. We can keep the issue open if it is still relevant. Please confirm if we need to keep the issue open.

Thank you for your support and cooperation.

[gaikwadrahul8]

Hi, @nsthorat

We have not received any confirmation from you, so we are closing this issue now. If the issue is still relevant, please let us know, either we'll re-open this issue or please feel free to create new issue after trying with latest version of Tensorflow.js. Thank you!