Размер модели меньше в формате .onnx, чем в формате .tflite
У меня есть предварительно обученная модель PyTorch, которую я хочу преобразовать в TFlite. Модель взята из API seisbench. Я использовал приведенный ниже код для преобразования. В коде есть несколько проверок, подтверждающих, что различные преобразования формата работают.
Я следовал за потоком .pt -> .onnx -> tensorflow -> tflite, но получил файл .onnx, который меньше (98 КБ), чем окончательная модель tflite (108 КБ). Я использую библиотеку onnx-tensorflow для преобразования файла .onnx в tensorflow ( )
model = sbm.PhaseNet.from_pretrained("instance") #load the model from the seisbench api
#model.load_state_dict(pNET.state_dict())
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# Save model information
print(model.get_model_args())
input_lenght = model.in_samples
input_depth = model.in_channels
# save to .pt
model.eval() #turn off gradient computations and other training-only operations
torch.save(model, 'pNET.pt')
# check if the model has been saved correctly
temp_model = torch.load('pNET.pt')
temp_model.eval()
print("Model's state_dict:")
for param_tensor in temp_model.state_dict():
print(param_tensor, "\t", temp_model.state_dict()[param_tensor].size())
# save to .onnx
# define an input vector (random vector)
sample_input = torch.randn(1, input_depth, input_lenght, requires_grad=True) #order is width, depth, lenght of input
#width fixed to 1 for time series data
# export
torch.onnx.export(
model, # PyTorch Model
sample_input, # Input tensor
'pNET.onnx', # Output file name
input_names=['input'], # Input tensor name (arbitrary)
output_names=['output'] # Output tensor name (arbitrary)
)
# check if the model has been saved correctly
onnx_model = onnx.load('pNET.onnx')
# Check that the IR is well formed
onnx.checker.check_model(onnx_model)
# Print a Human readable representation of the graph
onnx.helper.printable_graph(onnx_model.graph)
# Try to run an inference with the newly saved onnx model
import onnxruntime as ort
import numpy as np
ort_session = ort.InferenceSession('pNET.onnx')
outputs = ort_session.run(
None,
{'input': np.random.randn(1, input_depth, input_lenght).astype(np.float32)} #random input
)
print(outputs) #check if you get a tensor of the right shape
print(output_data.shape)
from onnx_tf.backend import prepare
# Converting to TensorFlow model
onnx_model = onnx.load("pNET.onnx") # load onnx model
tf_rep = prepare(onnx_model) # prepare tf representation
tf_rep.export_graph("pNET") # export the model
# Check if the conversion worked
# Run a TF inference
import tensorflow as tf
model = tf.saved_model.load("./pNET")
model.trainable = False
input_tensor = tf.random.uniform([1, input_depth, input_lenght])
out = model(**{'input': input_tensor})
print(out) #check if you get a tensor of the right shape
print(output_data.shape)
# float16 quantization
converter = tf.lite.TFLiteConverter.from_saved_model("./pNET")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_quant_model = converter.convert()
# Save the model
with open('pNETlite16float.tflite', 'wb') as f:
f.write(tflite_model) # same size as when I use interpreter instead of converter?
Мое замешательство связано с тем, что я ожидал, что квантование после обучения уменьшит размер модели. Добавляет ли TFLite какие-то дополнительные оболочки или методы в модель, увеличивая размер по сравнению с .onnx?