Proper way to save params with Relax detach_params

#1

I was able to get my model compiling using the reference on TVM’s documentation here:

However, part of the reference has the model and params split like so:

mod, params = relax.frontend.detach_params(mod)

The problem is that params is a dictionary, but not what something like tvm.runtime.save_param_dict(params) is expecting. Specifically, the expected type is Dict[str, NDArray], however the type of the params from the function is Dict[str, List[NDArray]].

Trying to save here fails with the below exception, since it is a list of NDArrays as the values, rather than just NDArray. What is the proper way to detach these save these params?

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/mitchell/micromamba/envs/mlc/lib/python3.11/site-packages/tvm/runtime/params.py", line 62, in save_param_dict
    return _ffi_api.SaveParams(_to_ndarray(params))
                               ^^^^^^^^^^^^^^^^^^^
  File "/home/mitchell/micromamba/envs/mlc/lib/python3.11/site-packages/tvm/runtime/params.py", line 27, in _to_ndarray
    transformed[k] = ndarray.array(v)
                     ^^^^^^^^^^^^^^^^
  File "/home/mitchell/micromamba/envs/mlc/lib/python3.11/site-packages/tvm/runtime/ndarray.py", line 675, in array
    return empty(arr.shape, arr.dtype, device, mem_scope).copyfrom(arr)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mitchell/micromamba/envs/mlc/lib/python3.11/site-packages/tvm/runtime/ndarray.py", line 430, in empty
    dtype = DataType(dtype)
            ^^^^^^^^^^^^^^^
  File "/home/mitchell/micromamba/envs/mlc/lib/python3.11/site-packages/tvm/_ffi/runtime_ctypes.py", line 182, in __init__
    raise ValueError("Do not know how to handle type %s" % type_str)
ValueError: Do not know how to handle type object
#2

the same error, did you solve this?

#3

I use this to solve the error:

def save_params(params, save_dir):
    if not isinstance(params, dict):
        raise ValueError("The params is not the dict")
    os.makedirs(save_dir, exist_ok=True)
    for index, param in enumerate(params["main"]):
        param_name = f"param_{index:03d}"
        param_path = os.path.join(save_dir, param_name+".bin")
        tvm.runtime.save_param_dict_to_file({param_name:param}, param_path)
        # with open(os.path.join(save_dir, param_name + ".bin"), "wb") as f:
        #     f.write(tvm.runtime.save_param_dict({param_name:param}))

def load_params(save_dir):
    if not os.path.exists(save_dir):
        raise ValueError(f"The params_dir does not exist: {save_dir}")
    
    params_value = []
    for param_file in sorted(os.listdir(save_dir)):  
        if param_file.endswith(".bin"):
            param_path = os.path.join(save_dir, param_file)
            param_value = tvm.runtime.load_param_dict_from_file(param_path)
            param_name = param_file.split('.')[-2]
            params_value.append(param_value[param_name])
    return {"main": params_value}