To avoid costly conversion from PyTorch, I have been saving the converted mod, params objects to file, as suggested by at least one of the tutorials.
with open(os.path.join(output_dir, relay_file), "w") as fo:
fo.write(tvm.ir.save_json(mod))
with open(os.path.join(output_dir, relay_params), "wb") as fo:
fo.write(tvm.runtime.save_param_dict(params))
This is a handy feature, however I have found that some of the info I need is not being saved properly. This means that for my desired feature to work, I have to regenerate mod, params from scratch each time.
My use-case is a bit of an edge-case, but definitely a growing part of core TVM. I’d be happy to try and fix this with a PR, if someone with more knowledge about the serialization format code could give some pointers.
context
For sparse computations, it is essential that we know that sparsity pattern of our data. Therefore, when we convert our model to sparse, code under relay/analysis converts and saves information about these buffers.
For example, in tvm/relay/analysis/sparse_conv2d.py:
...
register_task_input_buffer(
"default",
prefix + "W_data",
tvm.runtime.ndarray.array(sparse_weight_data),
overwrite=True,
)
register_task_input_buffer(
"default",
prefix + "W_indices",
tvm.runtime.ndarray.array(sparse_weight.indices),
overwrite=True,
)
...
This data is stored under the key default. Normally this isn’t used unless we want to auto-schedule. In that case, the auto-scheduler tries to load this data, in the function get_task_input_buffer() in auto_scheduler/search_task.py. In there we have the code:
...
# Then check for the default table, the input names extracted from a relay model will be
# stored here for we're not able to get the workload_key at that time
if input_name not in input_table:
input_table = TASK_INPUT_BUFFER_TABLE["default"]
...
However, if we save “sparsified” mod, params using the above methods, then try and save them, I have found that the "default" key is not saved, meaning that attempting to auto-schedule will crash on the input_table = TASK_INPUT_BUFFER_TABLE["default"] line.
Somehow the saving function is dropping the “default” key, but I’m not sure how or where. Any idea of how I can ensure it is saved correctly?
In theory, I could save the non-sparsified model, and sparsify after loading, but that doesn’t work either, since attempting to sparsify a mod, params loaded from file fails with:
File "/app/source/tvm/python/tvm/relay/analysis/sparse_conv2d.py", line 115, in process_params
del params[name]
TypeError: 'Map' object does not support item deletion
Evidently, the loaded mod, params do not match the saved ones.