Change code into Tensorflow2 APIs

[DanielMao1]

Dear deepmd-kit developers,
I am trying to improve the GPU training process. I found that when I am training, the training step spend most of time in 'kernel launch'. I did the profiling for several days, and I think the training process have much small kernels(or tensorflow operations) which can be fused by tensorflow XLA. But this process requires me use 'tf.function' API, which is used in tensorflow2.
I would like to have a try to change the code into tensorflow2 APIs which may enable me use XLA to do GPU kernel fusion, which I think can boosting the training step.
My question is: do you think too much small kernels is bottleneck in small data set(192 atoms, 40000 frames, datasize<200MB)? Is this method a proper way to improve training process?

[njzjz]

In fact, there is an api called tf.xla.experimental.jit_scope which is designed for operators, so I don't think it is necessary to switch to tf.function.

[DanielMao1]

Thank you. I will further explore this API. : )

[AbtMyML]

Hello! From your question, I learned that you also participated in the asc22 competition, so I want to ask you a question, do you use the source code or the deepmd-kit installed by conda command?
I've been puzzled and look forward to your reply. @DanielMao1

[DanielMao1]

yes. pip and conda

[AbtMyML]

Do you use conda to install the deepmd-kit? @DanielMao1

[DanielMao1]

no. Just pip install .