Added more detailed error reporting for post-training quantization
PiperOrigin-RevId: 272544323
diff --git a/tensorflow/lite/python/optimize/calibration_wrapper.cc b/tensorflow/lite/python/optimize/calibration_wrapper.cc
index 5da04f1..d2c2157 100644
--- a/tensorflow/lite/python/optimize/calibration_wrapper.cc
+++ b/tensorflow/lite/python/optimize/calibration_wrapper.cc
@@ -164,14 +164,19 @@
}
if (PyArray_NDIM(array) != tensor->dims->size) {
- PyErr_SetString(PyExc_ValueError, "Cannot set tensor: Dimension mismatch");
+ PyErr_Format(
+ PyExc_ValueError,
+ "Cannot set tensor: Dimension count mismatch, expected %d but found %d",
+ tensor->dims->size, PyArray_NDIM(array));
return nullptr;
}
for (int j = 0; j < PyArray_NDIM(array); j++) {
if (tensor->dims->data[j] != PyArray_SHAPE(array)[j]) {
- PyErr_SetString(PyExc_ValueError,
- "Cannot set tensor: Dimension mismatch");
+ PyErr_Format(PyExc_ValueError,
+ "Cannot set tensor: Size mismatch, expected %d for dim "
+ "%d but found %ld",
+ tensor->dims->data[j], j, PyArray_SHAPE(array)[j]);
return nullptr;
}
}
diff --git a/tensorflow/lite/python/optimize/calibrator_test.py b/tensorflow/lite/python/optimize/calibrator_test.py
index f027d6c..e94f0bc 100644
--- a/tensorflow/lite/python/optimize/calibrator_test.py
+++ b/tensorflow/lite/python/optimize/calibrator_test.py
@@ -124,7 +124,7 @@
for _ in range(10):
yield [np.ones(shape=(1, 2, 2, 3), dtype=np.float32)]
- with self.assertRaisesWithRegexpMatch(ValueError, 'Dimension mismatch'):
+ with self.assertRaisesWithRegexpMatch(ValueError, 'Size mismatch'):
quantizer.calibrate_and_quantize(input_gen, constants.FLOAT,
constants.FLOAT, False)
