Customize SI prefix in logging

python/tvm/autotvm/tuner/callback.py

@@ -23,6 +23,7 @@
 import numpy as np
 
 from .. import record
+from ..util import format_si_prefix
 
 logger = logging.getLogger('autotvm')
 
@@ -105,7 +106,7 @@ def trial_timestamps(self):
         return np.array(self.timestamps)
 
 
-def progress_bar(total, prefix=''):
+def progress_bar(total, prefix='', si_prefix='G'):
     """Display progress bar for tuning
 
     Parameters
@@ -114,6 +115,8 @@ def progress_bar(total, prefix=''):
         The total number of trials
     prefix: str
         The prefix of output message
+    si_prefix: str
+        SI prefix for flops
     """
     class _Context(object):
         """Context to store local variables"""
@@ -130,6 +133,9 @@ def __del__(self):
     ctx = _Context()
     tic = time.time()
 
+    # Validate si_prefix argument
+    format_si_prefix(0, si_prefix)
+
     if logger.level < logging.DEBUG:  # only print progress bar in non-debug mode
         sys.stdout.write('\r%s Current/Best: %7.2f/%7.2f GFLOPS | Progress: (%d/%d) '
                          '| %.2f s' % (prefix, 0, 0, 0, total, time.time() - tic))
@@ -147,10 +153,11 @@ def _callback(tuner, inputs, results):
             ctx.cur_flops = flops
             ctx.best_flops = tuner.best_flops
 
-            sys.stdout.write('\r%s Current/Best: %7.2f/%7.2f GFLOPS | Progress: (%d/%d) '
+            sys.stdout.write('\r%s Current/Best: %7.2f/%7.2f %sFLOPS | Progress: (%d/%d) '
                              '| %.2f s' %
-                             (prefix, ctx.cur_flops/1e9, ctx.best_flops/1e9, ctx.ct, ctx.total,
-                              time.time() - tic))
+                             (prefix, format_si_prefix(ctx.cur_flops, si_prefix),
+                              format_si_prefix(ctx.best_flops, si_prefix), si_prefix,
+                              ctx.ct, ctx.total, time.time() - tic))
             sys.stdout.flush()
 
     return _callback

python/tvm/autotvm/tuner/tuner.py

@@ -21,6 +21,7 @@
 import numpy as np
 
 from ..measure import MeasureInput, create_measure_batch
+from ..util import format_si_prefix
 
 from ..env import GLOBAL_SCOPE
 
@@ -87,7 +88,7 @@ def update(self, inputs, results):
         """
 
 
-    def tune(self, n_trial, measure_option, early_stopping=None, callbacks=()):
+    def tune(self, n_trial, measure_option, early_stopping=None, callbacks=(), si_prefix='G'):
         """Begin tuning
 
         Parameters
@@ -104,13 +105,18 @@ def tune(self, n_trial, measure_option, early_stopping=None, callbacks=()):
             (Tuner, List of MeasureInput, List of MeasureResult)
             with no return value. These callback functions will be called on
             every measurement pair. See autotvm/tuner/callback.py for some examples.
+        si_prefix: str
+            One of tvm.autotvm.util.SI_PREFIXES. The SI prefix to use when reporting FLOPS.
         """
         measure_batch = create_measure_batch(self.task, measure_option)
         n_parallel = getattr(measure_batch, 'n_parallel', 1)
         early_stopping = early_stopping or 1e9
         self.n_trial = n_trial
         self.early_stopping = early_stopping
 
+        # Validate si_prefix arg
+        format_si_prefix(0, si_prefix)
+
         old_level = logger.level
 
         GLOBAL_SCOPE.in_tuning = True
@@ -140,9 +146,9 @@ def tune(self, n_trial, measure_option, early_stopping=None, callbacks=()):
                     self.best_measure_pair = (inp, res)
                     self.best_iter = i + k
 
-                logger.debug("No: %d\tGFLOPS: %.2f/%.2f\tresult: %s\t%s",
-                             i + k + 1, flops / 1e9, self.best_flops / 1e9,
-                             res, config)
+                logger.debug("No: %d\t%sFLOPS: %.2f/%.2f\tresult: %s\t%s",
+                             i + k + 1, si_prefix, format_si_prefix(flops, si_prefix),
+                             format_si_prefix(self.best_flops, si_prefix), res, config)
 
             i += len(results)
             self.ttl = min(early_stopping + self.best_iter, n_trial) - i

python/tvm/autotvm/util.py

@@ -188,3 +188,12 @@ def get_const_tuple(in_tuple):
         else:
             ret.append(get_const_int(elem))
     return tuple(ret)
+
+
+SI_PREFIXES = 'yzafpn\xb5m kMGTPEZY'
+YOCTO_EXP10 = -24
+
+
+def format_si_prefix(x, si_prefix):
+    exp10 = 10 ** (SI_PREFIXES.index(si_prefix) * 3 + YOCTO_EXP10)
+    return float(x) / exp10

tests/python/unittest/test_format_si_prefix.py

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from numpy import isclose
+import random
+from tvm.autotvm import util
+
+
+SI_PREFIXES = 'yzafpn\xb5m kMGTPEZY'
+
+
+def test_format_si_prefix():
+  # test float conversion
+  assert util.format_si_prefix(1024, 'k') == 1.024
+
+  for i, prefix in enumerate(SI_PREFIXES):
+    integer, decimal = random.randint(0, 1000), random.randint(0, 1000)
+    exp = -24 + 3 * i   # 0th prefix (yocto) is 10^-24
+    number = integer * (10 ** exp) + decimal * (10 ** (exp - 3))
+    expected = (integer + decimal / 1000)
+    assert isclose(util.format_si_prefix(number, prefix), expected)
+
+  assert util.format_si_prefix(0, 'y') == 0
+
+
+if __name__ == '__main__':
+  test_format_si_prefix()