How to load PyTorch checkpoints into JAX/Flax?

[marcvanzee]

[marcvanzee]

Pytorch checkpoints contain a state_dict with all the weights/parameters for the models, and converting it to Flax involves:

1. Defining the model using Flax modules
2. Renaming the dictionary items to line up, and use the NCHW dimensions for conv weights.

Often flax.traverse_util.flatten_dict is useful, because you only need to operate on a flat dict instead of a nested dict. Once they align you use unflatten_dict to get the normal form back.

@nikitakit wrote the following code for importing PyTorch BERT checkpoints into a Flax model: https://github.com/nikitakit/flax_bert/blob/master/import_weights.py

[avital]

Here's another examples from Hugging Face BERT: https://github.com/huggingface/transformers/blob/a880f2549fd5652030afc244f3bb27ec764c5e43/src/transformers/models/bert/modeling_flax_bert.py#L452

[vaishnkv]

Hi, avital, I think this feature is not available as an API within the currently available transformers library? Is there any "official" way of doing this @avital ?

[vaishnkv]

This can be done with the "from_pt" argument (if we are loading a pre-trained model). I'm attaching screenshots for the same.

[davisyoshida]

@GCP20 There's no generic solution because people can choose to name their parameters different things between the two implementations. I wrote a helper script that does 90% of the work using a string similarity bipartite match, then I clean up the remainder manually.

[GinRawin]

Thank you for your guys' help. This method also works.