Efficient Ways for Saving and Loading weights

I’m sorry if it’s not the right place as I could not find the discussions or forum page.

I was wondering what are some of the most efficient ways to save and load models (also verify it’s properly loaded into GPU)?

1, In the docs, its given as save_the_model using Tensorflow 2. I also understand that the weights of haiku network are stored in a dictionary, as an example

{'linear': {'b': ndarray(..., shape=(300,), dtype=float32),
            'w': ndarray(..., shape=(28, 300), dtype=float32)},
 'linear_1': {'b': ndarray(..., shape=(100,), dtype=float32),
              'w': ndarray(..., shape=(1000, 100), dtype=float32)}}

Does haiku have some inbuilt function to save and load models? It becomes crucial in transfer learning tasks. Thanks in advance,

[tomhennigan]

Hi @VIGNESHinZONE , in general we recommend keeping a copy of your Python code which defines your model, and saving weights after training to disk. This provides the most flexibility. You can save weights produced by Haiku using pickle (https://colab.research.google.com/gist/tomhennigan/77a7e4ea04d716e9650bc54acf26f468/saving-and-loading-parameters-using-pickle.ipynb).

The integration with TensorFlow is an example of how to save the implementation of your model (e.g. the operations required to execute your neural network) as a TensorFlow graph, which is a suitable long term serialisation format for ML models. In general the best way to save the implementation of your model is to keep a copy of your Python code. In some situations (e.g. production serving) it is useful to have an intermediate format (like TF graphs) instead.

[tomhennigan]

Haiku itself doesn’t do anything special to support pmap/vmap. The key idea with Haiku is that once you’ve hk.transformed your function it is a “pure function” and you can use this directly with jax.pmap etc.

For an example of data parallel training with pmap see our resnet/imagenet example https://github.com/deepmind/dm-haiku/blob/master/examples/imagenet/train.py. In there the train_step function is parallel mapped across all available GPUs. You might observe that in this code we use Haiku to define a forward function with our neural net, and the rest is pure JAX with some other libraries in our ecosystem (e.g. optax for optimization).

As a meta point, its probably worth opening new issues for different questions to help other users who are searching.