I am trying to measure the change in size of PyTorch models after they are compiled with tvm relay.
I am comparing the file sizes after I write them on hard drive. Below is the code I use for this purpose:
def print_size_of_model(mod, label=" ", lib=None, is_tvm=False):
if is_tvm:
lib.export_library("temp.p")
size2 = os.path.getsize("temp.p")
size = size2
else:
torch.save(model.state_dict(), "temp.p")
size=os.path.getsize("temp.p")
print("model: ",label,' \t','Size (KB):', size/1e6)
os.remove('temp.p')
return size
# compare the sizes
f=print_size_of_model(model, "PyTorch")
q=print_size_of_model(mod=None, label = "tvm Relay", lib=lib, is_tvm=True)
# print("{0:.2f} times smaller".format(f/q))
md("### Result: {0:0.2f}x model size reduction by compilation with tvm relay".format(f/q))
the lib parameter that I use to represent the model generated from tvm.relay
is of type tvm.runtime.module.Module
. Maybe I need a better representation.
Result I get right now is
model: PyTorch | Size (KB): 422.502841 |
---|---|
model: tvm.relay | Size (KB): 1.189032 |
Result: 355.33x model size reduction by compilation with tvm relay
But I am afraid that it is not true. I appreciate any suggestion for correcting the result.