Merge pull request #83 from facebookresearch/tiny-improvement

Small improvement to parameters init fit autotuner

Author: prigoyal

tensor_comprehensions/tc_unit.py

@@ -159,39 +159,20 @@ def get_options_from_kwargs_and_tuner_cache(name, cache_file, options_cache, *in
 # TC autotuner class - ATen
 ###############################################################################
 class TcAutotuner(object):
-    def __init__(
-        self,
-        tc_lang,
-        pop_size=10,
-        crossover_rate=80,
-        mutation_rate=7,
-        generations=2,
-        number_elites=1,
-        threads=8,
-        gpus="0",
-        proto="/tmp/tuner.txt",
-        restore_from_proto=False,
-        restore_number=10,
-        log_generations=False,
-        tuner_min_launch_total_threads=64,
-        **kwargs
-    ):
+    def __init__(self, tc_lang, **kwargs):
         # tuner_cache will look like:
         # hash_key -> {"forward": options1, "backward": options2}
         self.tuner_cache = {}
         self.kwargs = kwargs
         self.tc_lang = tc_lang
         self.autotuner = ATenAutotuner(tc_lang)
-        self.set_autotuner_settings(
-            pop_size, crossover_rate, mutation_rate, generations, number_elites,
-            threads, gpus, proto, restore_from_proto, restore_number,
-            log_generations, tuner_min_launch_total_threads
-        )
+        self.set_autotuner_parameters(**kwargs)
 
-    def set_autotuner_settings(
-        self, pop_size, crossover_rate, mutation_rate, generations, number_elites,
-        threads, gpus, proto, restore_from_proto, restore_number, log_generations,
-        tuner_min_launch_total_threads,
+    def set_autotuner_parameters(
+        self, pop_size=10, crossover_rate=80, mutation_rate=7, generations=2,
+        number_elites=1, threads=8, gpus="0", proto="/tmp/tuner.txt",
+        restore_from_proto=False, restore_number=10, log_generations=False,
+        tuner_min_launch_total_threads=64, **kwargs
     ):
         self.autotuner.pop_size(pop_size)
         self.autotuner.crossover_rate(crossover_rate)
@@ -602,7 +583,7 @@ def autotune(self, *inputs, **kwargs):
         else:
             # we do the init again so that the autotuner parameters are updated
             # properly if users change them
-            self.tuner.__init__(self.lang, **kwargs)
+            self.tuner.set_autotuner_parameters(**kwargs)
         return self.tuner.autotune(*inputs, **kwargs)
 
 ###############################################################################

test_python/test_tc_torch.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 ##############################################################################
 
-import os, unittest, time, pdb
+import os, unittest, time
 
 import torch
 import torch.cuda
@@ -52,7 +52,6 @@ def convolution_grad(float(N,C,H,W) I, float(M,C,KH,KW) W1, float(N,M,H,W) O_gra
 }}
 """
 
-# PATH_PREFIX = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tc_test")
 PATH_PREFIX = os.path.join("/tmp/", "tc_test")
 
 if not os.path.exists(PATH_PREFIX):
@@ -333,7 +332,7 @@ def test_conv_train_autotune_no_cache_no_options_seed(self):
         I, W = torch.randn(N, C, H, W).cuda(), torch.randn(O, C, kH, kW).cuda()
         convolution.autotune(I, W, **tc.autotuner_settings)
         # on the second call, autotuning will be seeded from previous best options
-        convolution.autotune(I, W, **tc.autotuner_settings)
+        convolution.autotune(I, W, **tc.autotuner_settings, generations=5, pop_size=20)
 
     def test_conv_train_autotune_cache_no_options_seed(self):
         lang = CONV_TRAIN