Added switch --dont_save_keras_model to avoid writing model to disk

[travisjayday]

Problem

During TFLite conversion, openvino2tensorflow.py script calls

tf.lite.TFLiteConverter.from_keras_model(model) or tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)

Both these functions write or require weights/model to have been saved to disk. Saving a model to disk takes a long time because of I/O bottlenecks.

Solution

For backward compatibility, I added a switch --dont_save_keras_model that avoids saving the model to disk during TFLite conversion.

Implementation

Inspired by this SO post, we can create a TFLiteModel converter using concrete functions to avoid saving the model to disk. We call

tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])

instead of the above constructors.

Results

Model conversion time decreased from ~70s to ~25s for my particular model (MobileNetV3, 224x224x3), but time should relatively decrease more as model size increases.

[travisjayday]

@PINTO0309 do you have a model with multiple inputs/outputs? I'm trying to validate two input models with concrete_function tflite converter, but I'm having trouble getting such a model.

I created this simple OpenVino model model.zip that has two inputs and two outputs. However, openvino2tensorflow is throwing errors about dimension mismatch. Here is the code to generate the model:

from torch import nn

class TwoInputsNet(nn.Module):
  def __init__(self):
    super(TwoInputsNet, self).__init__()
    self.conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3)  # set up your layer here
    self.fc1 = nn.Linear(in_features=28, out_features=10 )  # set up first FC layer
    self.fc2 = nn.Linear(in_features=2868, out_features=10)  # set up the other FC layer

  def forward(self, input1, input2):
    c = self.conv(input1)
    f = self.fc1(input2)
    # now we can reshape `c` and `f` to 2D and concat them
    combined = torch.cat((c.view(c.size(0), -1),
                          f.view(f.size(0), -1)), dim=1)
    out = self.fc2(combined)
    out1 = self.fc2(combined)
    return out, out1

model = TwoInputsNet()
model.eval()
in1 = torch.rand((1, 3, 28, 28))
in2 = torch.rand((1, 3, 28, 28))

out = model(in1, in2)
torch.onnx.export(
    model,
    (in1, in2),
    './model/model.onnx',
    opset_version=13,
    input_names=['input_0', 'input_1'],
    output_names=['output_0', 'output_1'],
)
print(len(out))
print(out)

and here is the openvino2tensorflow error

ValueError: Dimensions must be equal, but are 3 and 28 for '{{node tf.linalg.matmul/MatMul}} = BatchMatMulV2[T=DT_FLOAT, adj_x=false, adj_y=true](Placeholder, tf.linalg.matmul/MatMul/b)' with input shapes: [1,28,28,3], [10,28].
ERROR: Please refer to 6-7 in the README first. https://github.com/PINTO0309/openvino2tensorflow

[PINTO0309]

There was insufficient consideration of the case where the MatMul is immediately after the 4D input layer. However, since the problem is essentially unrelated to the problem we want to solve with this pull request, it can be verified with a simpler multiplication-only model, as shown below.

model.zip

import torch
from torch import nn
import onnx
from onnxsim import simplify

class TwoInputsNet(nn.Module):
    def forward(self, input1, input2):
        a = input1 * 2.0
        b = input2 * 3.0

        out1 = a
        out2 = a + b

        return out1, out2

model = TwoInputsNet()
model.eval()
in1 = torch.rand((1, 3, 28, 28))
in2 = torch.rand((1, 3, 28, 28))

out = model(in1, in2)

onnx_file = 'model.onnx'
torch.onnx.export(
    model,
    (in1, in2),
    onnx_file,
    opset_version=11,
    input_names=['input_0', 'input_1'],
    output_names=['output_0', 'output_1'],
)
print(len(out))
print(out)

model_onnx2 = onnx.load(onnx_file)
model_simp, check = simplify(model_onnx2)
onnx.save(model_simp, onnx_file)

"""
mo \
--framework onnx \
--input_model model.onnx \
--data_type FP32 \
--output_dir openvino/FP32 \
--model_name model

openvino2tensorflow \
--model_path openvino/FP32/model.xml \
--output_saved_model \
--output_pb \
--non_verbose \
--output_no_quant_float32_tflite
"""

[travisjayday]

OK @PINTO0309, I made the changes you suggested. Multiple inputs are handled correctly now with try/catch logic in case concrete function fails. What do you think?

[PINTO0309]

LGTM. Thanks for your contribution. :+1: