geoelements/gns

Comparing results with Deepmind implementation?

Closed this issue · 3 comments

Hello,

Thanks so much for this repository. It is so valuable to have a Pytorch implementation of "Learning to simulate", which can really accelerate research progress in this field. Really appreciate this work.

I wanted to ask whether you have compared results with Deepmind's implementation that you've reproduced. In particular, it seems feasible to take some dataset (e.g. WaterRamps), train the model, and then evaluate the MSE on the test set, either 1-step or over a whole rollout. Have you done this and if so, seen comparable performance to the Deepmind implementation? I wanted to ask since it would be useful to compare performance. Thanks a lot.

kks32 commented

Hi @arjun-mani Thanks for your interest in our PyTorch implementation. We just finished training for 20 Million steps as done in the DeepMind's paper. The MSE during training were in the same range and the viz looks reasonable. We will post the comparison of MSE for rollouts and 1-step soon. Will add some performance comparisons as well. Thanks!

Thanks very much! Glad to hear, look forward to seeing the comparisons.

kks32 commented

@arjun-mani After 1 Million training steps - the RMSE is roughly similar. The training steps are randomized, so we don't expect exactly the same RMSE.

RMSE TF PyTorch
Sand 0.000947 0.01352

rollout_tf_sand_1M_ex0

TensorFlow

rollout_sand_1M_ex0

PyTorch