Размер модели меньше в формате .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?

0 ответов

Другие вопросы по тегам