JAX numpy预训练权重转H5格式报错，求正确转换方法

你遇到的错误核心原因是：直接将NpzFile容器对象传给h5py的create_dataset方法。.npz文件本质是多个numpy/JAX数组的压缩集合，并非单个数组，h5py无法直接识别这种容器类型，报错里的<U71字符串类型，就是误将npz的键名当成了数据导致的。

正确的转换代码 #

修改代码，遍历npz文件中的所有权重数组，逐个写入H5文件，同时使用下载得到的正确路径，避免硬编码文件名：

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}"

# 权重文件保存到本地路径"~/.keras/weights/"
local_filepath = tf.keras.utils.get_file(fname, origin, cache_subdir="weights")

# 加载npz并转换为h5
with jnp.load(local_filepath) as jax_weights, h5py.File('ViT-B_16_imagenet21k.h5', 'w') as hf:
    # 遍历每个权重数组，按原键名存入H5
    for key, array in jax_weights.items():
        # 将JAX数组转为numpy数组，确保h5py兼容
        np_array = array.__array__()
        hf.create_dataset(key, data=np_array)

关键说明 #

1. 使用with语句管理文件：同时处理npz和H5文件，自动完成资源关闭，避免文件泄漏
2. 遍历权重键值对：npz文件是键值对结构，遍历items()可以保留原权重的命名，方便后续在Keras等框架中加载
3. JAX数组转numpy：h5py对原生numpy数组支持更好，通过__array__()方法完成转换，不丢失数据
4. 使用下载的真实路径：用local_filepath替代硬编码的文件名，确保加载的是正确路径下的下载文件

内容的提问来源于stack exchange，提问作者MediaJ