Prepare to call print_unique() or analog in normalize.py

felker · felker · commit 55639ce76331 · 2019-11-05T15:30:39.000-06:00
diff --git a/examples/mpi_learn.py b/examples/mpi_learn.py
@@ -2,11 +2,9 @@
     mpi_train, mpi_make_predictions_and_evaluate,
     comm, task_index, print_unique
     )
-# from mpi4py import MPI
 from plasma.preprocessor.preprocess import guarantee_preprocessed
 from plasma.models.loader import Loader
 from plasma.conf import conf
-# from pprint import pprint
 '''
 #########################################################
 This file trains a deep learning model to predict
@@ -57,24 +55,11 @@
     print('unkown normalizer. exiting')
     exit(1)
 
-# TODO(KGF): this part of the code is duplicated in mpi_runner.py
-# comm = MPI.COMM_WORLD
-# task_index = comm.Get_rank()
-# num_workers = comm.Get_size()
-
-# NUM_GPUS = conf['num_gpus']
-# MY_GPU = task_index % NUM_GPUS
-# backend = conf['model']['backend']
-
-# if task_index == 0:
-#     pprint(conf)
-
 # TODO(KGF): confirm that this second PRNG seed setting is not needed
 # (before normalization; done again before MPI training)
 # np.random.seed(task_index)
 # random.seed(task_index)
 
-
 only_predict = len(sys.argv) > 1
 custom_path = None
 if only_predict:
@@ -94,8 +79,7 @@
 (shot_list_train, shot_list_validate,
  shot_list_test) = guarantee_preprocessed(conf)
 
-
-print_unique("begin normalization...")  # , end='')
+print_unique("begin normalization...")
 normalizer = Normalizer(conf)
 normalizer.train()
 loader = Loader(conf, normalizer)
diff --git a/plasma/preprocessor/normalize.py b/plasma/preprocessor/normalize.py
@@ -73,9 +73,9 @@ def save_stats(self):
     def load_stats(self):
         pass
 
-    @abc.abstractmethod
     def print_summary(self, action='loaded'):
-        pass
+        print('{} normalization data from {} shots ( {} disruptive )'.format(
+            action, self.num_processed, self.num_disruptive))
 
     def set_inference_mode(self, val):
         self.inference_mode = val
@@ -132,8 +132,7 @@ def train_on_files(self, shot_files, use_shots, all_machines):
             start_time = time.time()
 
             for (i, stats) in enumerate(pool.imap_unordered(
-                    self.train_on_single_shot,
-                    shot_list_picked)):
+                    self.train_on_single_shot, shot_list_picked)):
                 # for (i,stats) in
                 # enumerate(map(self.train_on_single_shot,shot_list_picked)):
                 if stats.machine in machines_to_compute:
@@ -308,10 +307,6 @@ def load_stats(self):
             print('Machine {}:'.format(machine))
             self.print_summary()
 
-    def print_summary(self, action='loaded'):
-        print('{} normalization data from {} shots ( {} disruptive )'.format(
-            action, self.num_processed, self.num_disruptive))
-
 
 class VarNormalizer(MeanVarNormalizer):
     def apply(self, shot):
@@ -341,7 +336,6 @@ def __str__(self):
 
 
 class AveragingVarNormalizer(VarNormalizer):
-
     def apply(self, shot):
         apply_positivity(shot)
         super(AveragingVarNormalizer, self).apply(shot)
@@ -444,17 +438,10 @@ def save_stats(self):
         # num_processed = dat['num_processed']
         # num_disruptive = dat['num_disruptive']
         self.ensure_save_directory()
-        np.savez(
-            self.path,
-            minimums=self.minimums,
-            maximums=self.maximums,
-            num_processed=self.num_processed,
-            num_disruptive=self.num_disruptive,
-            machines=self.machines)
-        print(
-            'saved normalization data from {} shots ( {} disruptive )'.format(
-                self.num_processed,
-                self.num_disruptive))
+        np.savez(self.path, minimums=self.minimums, maximums=self.maximums,
+                 num_processed=self.num_processed,
+                 num_disruptive=self.num_disruptive, machines=self.machines)
+        self.print_summary(action='saved')
 
     def load_stats(self):
         assert(self.previously_saved_stats()[0])
@@ -464,10 +451,9 @@ def load_stats(self):
         self.num_processed = dat['num_processed'][()]
         self.num_disruptive = dat['num_disruptive'][()]
         self.machines = dat['machines'][()]
-        print(
-            'loaded normalization data from {} shots ( {} disruptive )'.format(
-                self.num_processed,
-                self.num_disruptive))
+        for machine in self.means:
+            print('Machine {}:'.format(machine))
+            self.print_summary()
 
 
 def get_individual_shot_file(prepath, shot_num, ext='.txt'):
