larq / compute-engine

Highly optimized inference engine for Binarized Neural Networks
https://docs.larq.dev/compute-engine
Apache License 2.0
242 stars 34 forks source link

Add SetBatchSize pass to MLIR converter #690

Closed Tombana closed 2 years ago

Tombana commented 2 years ago

What do these changes do?

This adds a pass to the converter that sets the dynamic batchsize to 1 on the input layer.

When batch_size is not explicitly set to 1, it will remain a ? (wildcard) throughout all MLIR passes. Only at the final MLIR to Flatbuffer conversion stage, are all the ? simply converted to 1: https://github.com/tensorflow/tensorflow/blob/v2.7.0/tensorflow/compiler/mlir/lite/flatbuffer_export.cc#L844

This causes broken Reshape and Concatenation ops: the MLIR passes will correctly infer the Reshape output shape to something like ?x?x2, and then the flatbuffer exporter simply overrides both with 1. Other artifacts

The graphdef converter code has a batch_size=1 fix in Python, but the saved model converter did not have something like this yet.

How Has This Been Tested?

MLIR filecheck tests have been added.