Loading model from the saved data

[loshmii]

Hi I'm having trouble loading model from the data I have saved when training the model using orbax.checkpoint

ckpt_dir = Path('checkpoints').resolve()
ckpt_dir = ocp.test_utils.erase_and_create_empty(ckpt_dir)

_, state = nnx.split(model)
nnx.display(state)

checkpointer = ocp.StandardCheckpointer()
checkpointer.save(ckpt_dir / 'state', state)

The model is instance of the following class :

class Transformer(nnx.Module) :
  def __init__(self, rngs : nnx.Rngs, d_model, n_heads, n_layers, padding_id = 100257) :
    super().__init__()
    self.token_embedding = TokenEmbeddings(rngs=rngs, vocab_size=vocab_size, d_model=d_model)
    self.pos_enc = PositionalEncoding(rngs=rngs)
    self.transformer_blocks = [TransformerBlock(rngs=rngs, d_model=d_model, n_heads=n_heads) for _ in range(n_layers)]
    self.transformer_output = TransformerOutput(rngs=rngs, d_model=d_model, vocab_size=vocab_size)
    self.padding_id = padding_id

  def __call__(self, x : jax.Array) -> jax.Array :
    """Forward pass of the transformer model.

    Args :
      x : jax.Array
        The input tensor of shape [batch, seq_len]

    Return  :
      output : jax.Array
        The output tensor of shape [batch, seq_len, d_model]
    """

    padding_mask = jnp.where(x == self.padding_id, 1, 0)
    x = self.token_embedding(x)
    x = jnp.add(x, self.pos_enc(x))
    for transformer_block in self.transformer_blocks :
      x = transformer_block(x, padding_mask)
    return self.transformer_output(x)

The loading part of the code looks like :

abstract_model = nnx.eval_shape(lambda : Transformer(rngs=nnx.Rngs(0), d_model=d_model, n_heads=n_heads,n_layers=n_layers))
graphdef, abstract_state = nnx.split(abstract_model)
print('The abstract NNX state (all leaves are abstract arrays) :')
nnx.display(abstract_state)

checkpointer = ocp.StandardCheckpointer()
ckpt_dir = Path('checkpoints/state').resolve()
state_restored = checkpointer.restore(ckpt_dir, abstract_state)
jax.tree.map(np.testing.assert_array_equal, state, state_restored)
print('NNX State restored : ')
nnx.display(state_restored)

ChatGPT says there is a problem with RNG class but honestly I have no idea what's wrong with this way of loading the parameters. Everything except loading parameters works fine. I would greatly appreciate any help. The error I'm getting is :

TypeError                                 Traceback (most recent call last)
Cell In[24], [line 8](vscode-notebook-cell:?execution_count=24&line=8)
      [6](vscode-notebook-cell:?execution_count=24&line=6) checkpointer = ocp.StandardCheckpointer()
      [7](vscode-notebook-cell:?execution_count=24&line=7) ckpt_dir = Path('checkpoints/state').resolve()
----> [8](vscode-notebook-cell:?execution_count=24&line=8) state_restored = checkpointer.restore(ckpt_dir, abstract_state)
      [9](vscode-notebook-cell:?execution_count=24&line=9) jax.tree.map(np.testing.assert_array_equal, state, state_restored)
     [10](vscode-notebook-cell:?execution_count=24&line=10) print('NNX State restored : ')

File /venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:152, in StandardCheckpointer.restore(self, directory, target, strict)
    [125](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:125) def restore(
    [126](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:126)     self,
    [127](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:127)     directory: epath.PathLike,
   (...)
    [130](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:130)     strict: bool = True,
    [131](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:131) ) -> Any:
    [132](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:132)   """Restores a checkpoint.
    [133](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:133) 
    [134](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:134)   Args:
   (...)
    [150](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:150)     The restored checkpoint.
    [151](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:151)   """
--> [152](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:152)   return super().restore(
    [153](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:153)       directory, args=StandardRestore(target, strict=strict)
    [154](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py:154)   )
...
--> [494](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:494)       'dtype': jnp.dtype(args.dtype).name,
    [495](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:495)   }
    [496](https://vscode-remote+ssh-002dremote-002b142-002e113-002e12-002e8.vscode-resource.vscode-cdn.net/venv/main/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:496) return tspec

TypeError: Cannot interpret 'key<fry>' as a data type

[loshmii]

I have no idea what happened with formatting, sorry everyone.

[jakevdp]

I fixed it – you should use code blocks next time: https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks

[jakevdp]

This function looks specific to flax/nnx. You might have more luck asking at https://github.com/google/flax

[loshmii]

For anyone having the same problem I was able to find the solution - google/orbax#1105