Added a check for viable quantization layers to annotation.

apache · May 14, 2019 · f06693e · f06693e
1 parent bbf128a
commit f06693e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/tvm/relay/quantize/_annotate.py b/python/tvm/relay/quantize/_annotate.py
@@ -163,6 +163,11 @@ def conv2d_rewrite(ref_call, new_args, ctx):
             _set_conv_counter(cnt + 1)
             return None
 
+    # Check if the kernel is suitable for quantization, post warning if not.
+    in_channels = new_args[1].data.shape[1]
+    if in_channels % 4 != 0:
+        print("Conv layer %d may not be suitable for quantization. Consider leaving it FP." %cnt)
+
     _set_conv_counter(cnt + 1)
 
     lhs_expr, lhs_kind = _get_expr_kind(new_args[0])