Batches handled differently in the int reference fully connected kernel

[mansnils]

System information
-- TensorFlow installed from (source or binary):

• TensorFlow version (use command below): SHA1: 403b4b1

Describe the current behavior
Almost all fully connected kernels handle batches differently than the int8 reference fully connected kernel. Is this really intentional or a bug?

We have seen models (with keep_num_dims=True for fully_connected) where batches can be different in fully connected when a model has been converted from float to int8.

find . -name "fully_connected.h" |xargs grep "const int batches"|grep -v "//"
./lite/kernels/internal/reference/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);
./lite/kernels/internal/reference/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
./lite/kernels/internal/reference/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
./lite/kernels/internal/reference/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
./lite/kernels/internal/reference/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
./lite/kernels/internal/reference/integer_ops/fully_connected.h: const int batches = output_shape.Dims(0);
./lite/kernels/internal/reference/integer_ops/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
./lite/kernels/internal/optimized/sparse_ops/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);
./lite/kernels/internal/optimized/sparse_ops/fully_connected.h: const int batches = thread_end - thread_start;
./lite/kernels/internal/optimized/sparse_ops/fully_connected.h: const int batches =
./lite/kernels/internal/optimized/integer_ops/fully_connected.h: const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);

The typical case is const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1)
In the int8 ref case: const int batches = output_shape.Dims(0);

Describe the expected behavior
Batches are handled the same way for int and float fc kernels.

Contributing

• Do you want to contribute a PR? (yes/no): yes
• Briefly describe your candidate solution(if contributing):
  Change const int batches = output_shape.Dims(0);
  To
  const int output_dims_count = output_shape.DimensionsCount();
  const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);

[dansitu]

We are also seeing this issue. As a workaround, conv layers with 1x1 filter kernels can substitute for fully connected layers. Note that for quantized models this will result in per-channel quantization being applied, which may impact your model size and output.

There's a related issue here tracking fixing the problem for optimized kernels on Arm embedded devices that use TensorFlow Lite for Microcontrollers: ARM-software/CMSIS_5#1398

[Tessil]

With #54223 merged the problem should now be solved. The fix was also imported in tflite-micro (see commit tensorflow/tflite-micro@1501b57).

[sachinprasadhs]

@mansnils , Could you please try it in Nightly version and let us know if the problem is fixed as per the above comment. Thanks!

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please reopen if you'd like to work on this further.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No