Error through converting a jax numpy pre-trained weight to h5 weight

[JavanehBahrami]

Hi all.
I have downloaded a jax numpy weight file with npz suffix, but when I tried to convert it to h5 file I recieved this error:

import jax.numpy as jnp
import h5py
import tensorflow as tf


BASE_URL = "https://github.com/faustomorales/vit-keras/releases/download/dl"

size = "B_16"
weights = "imagenet21k"
fname = f"ViT-{size}_{weights}.npz"
origin = f"{BASE_URL}/{fname}"

# saved weight file in local path "~/.keras/weights/"
local_filepath = tf.keras.utils.get_file(fname, origin, cache_subdir="weights")

# Load the npz and convert to h5
jax_file = jnp.load('ViT-B_16_imagenet21k.npz')
with h5py.File('ViT-B_16_imagenet21k.h5', 'w') as hf:
    hf.create_dataset('weights', data=jax_file)

the npz file will be downloaded in local path "~/.keras/weights/".

my error is:

*** TypeError: No conversion path for dtype: dtype('<U71')
Traceback (most recent call last):
  File "/home/javaneh/jb_env/lib/python3.8/site-packages/h5py/_hl/group.py", line 161, in create_dataset
    dsid = dataset.make_new_dset(group, shape, dtype, data, name, **kwds)
  File "/home/javaneh/jb_env/lib/python3.8/site-packages/h5py/_hl/dataset.py", line 88, in make_new_dset
    tid = h5t.py_create(dtype, logical=1)
  File "h5py/h5t.pyx", line 1663, in h5py.h5t.py_create
  File "h5py/h5t.pyx", line 1687, in h5py.h5t.py_create
  File "h5py/h5t.pyx", line 1753, in h5py.h5t.py_create

my question is, How can I convert a jax file into h5 file, correctly?

Note1: the output of dir(jax_file):

dir(jax_file)
['__abstractmethods__', '__class__', '__contains__', '__del__', '__delattr__', '__dict__', '__dir__', '__doc__', '__enter__', '__eq__', '__exit__', '__format__', '__ge__', '__getattribute__', '__getitem__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__iter__', '__le__', '__len__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__reversed__', '__setattr__', '__sizeof__', '__slots__', '__str__', '__subclasshook__', '__weakref__', '_abc_impl', '_files', 'allow_pickle', 'close', 'f', 'fid', 'files', 'get', 'items', 'keys', 'pickle_kwargs', 'values', 'zip']

Note2: type of jax_file is:

type(jax_file)
<class 'numpy.lib.npyio.NpzFile'>

Note3: my tensorflow version 2.9.1

Any help will be appreciated.

[jakevdp]

Hi - thanks for the question! This actually has nothing to do with JAX: jnp.savez is a convenience wrapper that converts JAX arrays to numpy arrays and stores them via numpy in a .npz file, and at that point JAX is out of the picture.

It looks like the problem you're facing comes from misuse of the h5py create_dataset function. According to the docs, the data argument must be a numpy array. An npz file object is not a numpy array, but a view of an on-disk collection of numpy arrays.

So, for example, if you want to create your h5 dataset from the first array in the file, you could do something like this:

jax_file = jnp.load('ViT-B_16_imagenet21k.npz')
arr = jax_file[jax_file.files[0]]
with h5py.File('ViT-B_16_imagenet21k.h5', 'w') as hf:
  hf.create_dataset('weights', data=arr)

You can read more about npz files and how to interact with them in the numpy docs.

[JavanehBahrami]

Thanks @jakevdp for replying,
the challenge is that, it should be a unified h5 file(not multiple one). because I need to load h5 file once and feed to my tensorflow model.
length of jax_file.files is more 200 items, but how to unified all 200 items in one weight file in h5 format?

is there any way to create a unified h5 file from jax numpy file?

[jakevdp]

Got it, thanks. Just to be clear, this question has nothing to do with JAX (npz files are a numpy concept and contain numpy arrays), and I've never used h5py. But looking at the h5py docs, it emphasizes this:

Groups work like dictionaries, and datasets work like NumPy arrays

You were calling create_dataset, which expects a numpy array. It seems like an npz file (a collection of numpy arrays) conceptually maps to an h5py group, not an h5py dataset.

Further down on the quickstart page, it shows examples of how to create a group:

>>> f = h5py.File('mydataset.hdf5', 'a')
>>> grp = f.create_group("subgroup")
>>> dset2 = grp.create_dataset("another_dataset", (50,), dtype='f')

So I suspect what you'll need to do to create a group from an npz file is to use group = f.create_group(), and then call group.create_dataset() once per array in your npz file. Maybe something like this:

with h5py.File('ViT-B_16_imagenet21k.h5', 'w') as hf:
  group = hf.create_group('mydata'):
  for name, arr in jax_file.items():
    group.create_dataset(name, data=arr)

If that's not what you're looking for, then I'd suggest asking in h5py-specific help forums. Best of luck!