Splitting an ONNX DNN Model

Question

I'm trying to split DNN Models in order to execute part of the network on the edge and the rest on the cloud. Because it has to be cross-platform and work with every framework I need to do it directly starting from an ONNX model.

I know how to generate an ONNX model starting from tensorflow/keras and how to run an ONNX model, but I realized that is really hard to work on the ONNX file, like visualizing it and modify it.

Is there someone that can help me understand how to split and ONNX model, or at least run part of an ONNX model (like from input to layer N and from layer N to the output)?

I'm starting from this situation:

# load MobileNetV2 model
model = MobileNetV2()

# Export the model
tf.saved_model.save(model, "saved_model")  

# export to .onnx
!python -m tf2onnx.convert --saved-model saved_model --output mobilenet_v2.onnx --opset 7

# open the saved ONNX Model
print("Import ONNX Model..")
onnx_model = onnx.load("mobilenet_v2.onnx")
tf_rep = prepare(onnx_model, logging_level="WARN", auto_cast=True)

I tried to use sclblonnx but on models this big(although it's a small model) I can't really print the graph and when I list the inputs and outputs with textlist_inputs/list_outputs I don't really get how ther are interconnected.

Any help would be greatly appreciated. Thank you in advance.

Answer 1

From Onnx PythonAPI specs, you can split onnx model by specifying input name and output name of the tensors.