Setting `ENGINE_COMPILE_ONLY` to `true` is resulting in output tensor of `0s`

[mohit-av]

I have a tfjs model which is taking an inference time of 1.1 sec when I'm passing a batch of input. I came across this which suggest that using this will change compilation mode from sequence to parallel. It will result in the the lower inference time.

So, I added this into my code:

const t = Date.now();
tf.tidy(() => {
  tf.env().set('ENGINE_COMPILE_ONLY', true);
  console.log("warm up output",model.predict({ "input_2:0": tf.tensor(arr) }).dataSync());
  tf.backend().checkCompileCompletion();
  tf.backend().getUniformLocations();
  tf.env().set('ENGINE_COMPILE_ONLY', false);
});
console.log("warm up time", Date.now() - t);

It decreased my inference time to 120 millisec. But now, I'm getting tensor of 0s as output. It seems changing all the values to 0s is the reason for lower inference time.

I have set tf.setBackend('webgl') and I'm using tfjs@4.2.0.

My model architecture looks like this:
So, I take the x,y co-ordinates as input. Then I'm using positional encoding on this 2D data to transform to higher space of 42. Next there are 5 dense layer. The first layer input layer with 42 neurons with sine activation, then next 3 layer with 256 neurons each with sine activation and then final layer with 3 neurons.

Could anyone please suggest what's wrong in this? or any other way to decrease my inference time?

[qjia7]

ENGINE_COMPILE_ONLY is only for compile. So you need to call model.predict again for the real execution after the compile pass.

[mohit-av]

@qjia7, If I added it again after setting the environment variable back to false to allow the real execution. But then it will result in the inference time of 1.2 sec (1st inference + 2nd inference). How can I use it to boost my performance? Please correct me if I'm wrong.

const t = Date.now();
tf.tidy(() => {
  tf.env().set('ENGINE_COMPILE_ONLY', true);
  console.log("warm up output",model1.predict({ "input_2:0": tf.tensor(arr) }).dataSync());  #---> 1st inference 
  tf.backend().checkCompileCompletion();
  tf.backend().getUniformLocations();
  tf.env().set('ENGINE_COMPILE_ONLY', false);
  console.log("warm up output",model1.predict({ "input_2:0": tf.tensor(arr) }).dataSync());   #---> 2nd inference 
});
console.log("warm up time", Date.now() - t);

[qjia7]

Do you see any difference between 1 and 2? The browser itself has the cache. If you run your model once, the browser will cache the shader binary and the ENGINE_COMPILE_ONLY will be not that necessary. But if you first run your model or clear the browser cache, you may see ENGINE_COMPILE_ONLY greatly speed up the shader compilation (the 1st warmup). What's your model's average predict time by 100 times (maybe you can exclude the first several times to get a more stable data) to exclude the model itself is very slow?
1.

const t = Date.now();
tf.tidy(() => {
  tf.env().set('ENGINE_COMPILE_ONLY', true);
  model1.predict({ "input_2:0": tf.tensor(arr) }).dataSync();
  tf.backend().checkCompileCompletion();
  tf.backend().getUniformLocations();
  tf.env().set('ENGINE_COMPILE_ONLY', false);
  model1.predict({ "input_2:0": tf.tensor(arr) }).dataSync();
});
console.log("warm up time", Date.now() - t);

2

const t = Date.now();
tf.tidy(() => {
  model1.predict({ "input_2:0": tf.tensor(arr) }).dataSync();
  model1.predict({ "input_2:0": tf.tensor(arr) }).dataSync(); 
});
console.log("warm up time", Date.now() - t);

[mohit-av]

@qjia7, So, you mean that we are supposed to use this only for reducing warm up time. I have a good understanding that we need to warm up the model first with some test data, so all the shaders are compiled and cached on the GPU. But what should I do to decrease inference time?

P.S.: I have already tried quantization, pruning, weight clustering but they only reduce the model size but no affect on inference time.

[qjia7]

So, you mean that we are supposed to use this only for reducing warm up time

Yes. This flag has no effect for inference time.

what should I do to decrease inference time?

If you can share your model, we can help debug what's the bottleneck. Or you can profile your model to see which ops with what kind of input shapes are consuming lots of time and report bug to us. For example, https://tfjs-benchmarks.web.app/local-benchmark/, you can choose kernelTiming with aggregate/individual mode and click Run benchmarks to see the detailed per op execution time to find out culprit. In your case, you can choose the mode custom and paste the model url and see what you can find out.

And you can also try webgpu backend and let us know the result:). webgpu is very friendly for middle/large models.

[qjia7]

Forgot one thing. Maybe for webgl backend, you can try tf.env().set('WEBGL_EXP_CONV', true); before predicting to see if you can get better performance.

cc @Linchenn

[mohit-av]

@qjia7, Thanks for sharing this. I'm not using any conv layer. So, I guess tf.env().set('WEBGL_EXP_CONV', true); won't work for me. Although, I'm using this tf.env().set('WEBGL_USE_SHAPES_UNIFORMS', true); which has no affect on my inference time.
Here is the link for my model:

1. model.json
2. group1-shared1of1.bin

Its input shape is [-1, 2] and output shape is [-1,3]. I have set the batch size to -1 so, i could send any shape input. For your reference, I'm testing on a batch size of 262144 i.e. for input shape: [262144, 2]

[mohit-av]

So, I ran it on https://tfjs-benchmarks.web.app/local-benchmark/. For me, these ops are taking more time:

1. _FusedMatMul ,
2. sin,
3. fusedMatMul__op

I'm using 4 dense layer with sin activation which are using using these ops. Is there any way to replace these ops with the ones which are faster without sacrificing the accuracy.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No