profiling wip

baselines/common/running_mean_std.py

@@ -52,38 +52,25 @@ class TfRunningMeanStd(object):
             self._var   = tf.get_variable('std',   initializer=np.ones(shape, 'float64'),       dtype=tf.float64)    
             self._count = tf.get_variable('count', initializer=np.full((), epsilon, 'float64'), dtype=tf.float64)
 
-        self.update_ops = [
+        self.update_ops = tf.group([
             self._var.assign(self._new_var),
             self._mean.assign(self._new_mean),
             self._count.assign(self._new_count)
-        ]
+        ])
 
         sess.run(tf.variables_initializer([self._mean, self._var, self._count]))
         self.sess = sess
-
+        self._set_mean_var_count()
     
-                     
-    @property
-    def mean(self):
-        return self.sess.run(self._mean)
-
-    @property
-    def var(self):
-        return self.sess.run(self._var)
-
-    @property
-    def count(self):
-        return self.sess.run(self._count)
-
+    def _set_mean_var_count(self):
+        self.mean, self.var, self.count = self.sess.run([self._mean, self._var, self._count])                    
          
     def update(self, x):
         batch_mean = np.mean(x, axis=0)
         batch_var = np.var(x, axis=0)
         batch_count = x.shape[0]
 
-        mean, var, count = self.sess.run([self._mean, self._var, self._count])
-        new_mean, new_var, new_count = update_mean_var_count_from_moments(mean, var, count, batch_mean, batch_var, batch_count)
-
+        new_mean, new_var, new_count = update_mean_var_count_from_moments(self.mean, self.var, self.count, batch_mean, batch_var, batch_count)
 
         self.sess.run(self.update_ops, feed_dict={
             self._new_mean: new_mean,
@@ -91,6 +78,8 @@ class TfRunningMeanStd(object):
             self._new_count: new_count
         })
 
+        self._set_mean_var_count()
+
         
 
 def test_runningmeanstd():
@@ -126,3 +115,71 @@ def test_tf_runningmeanstd():
         ms2 = [rms.mean, rms.var]
 
         np.testing.assert_allclose(ms1, ms2)
+
+
+def profile_tf_runningmeanstd():
+    import time
+    from baselines.common import tf_util
+
+    tf_util.get_session( config=tf.ConfigProto(
+        inter_op_parallelism_threads=1,
+        intra_op_parallelism_threads=1,
+        allow_soft_placement=True
+    ))
+
+    x = np.random.random((376,))
+
+    n_trials = 10000
+    rms = RunningMeanStd()
+    tfrms = TfRunningMeanStd()
+
+    tic1 = time.time()
+    for _ in range(n_trials):
+        rms.update(x)
+
+    tic2 = time.time()
+    for _ in range(n_trials):
+        tfrms.update(x)
+
+    tic3 = time.time()
+
+    print('rms update time ({} trials): {} s'.format(n_trials, tic2 - tic1))
+    print('tfrms update time ({} trials): {} s'.format(n_trials, tic3 - tic2))
+    
+
+    tic1 = time.time()
+    for _ in range(n_trials):
+        z1 = rms.mean
+
+    tic2 = time.time()
+    for _ in range(n_trials):
+        z2 = tfrms.mean
+
+    tic3 = time.time()
+
+    print('rms get mean time ({} trials): {} s'.format(n_trials, tic2 - tic1))
+    print('tfrms get mean time ({} trials): {} s'.format(n_trials, tic3 - tic2))
+         
+    
+
+    '''
+    options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) #pylint: disable=E1101
+    run_metadata = tf.RunMetadata()
+    profile_opts = dict(options=options, run_metadata=run_metadata)
+
+    
+
+    from tensorflow.python.client import timeline
+    fetched_timeline = timeline.Timeline(run_metadata.step_stats) #pylint: disable=E1101
+    chrome_trace = fetched_timeline.generate_chrome_trace_format()
+    outfile = '/tmp/timeline.json'
+    with open(outfile, 'wt') as f: 
+        f.write(chrome_trace)
+    print(f'Successfully saved profile to {outfile}. Exiting.')
+    exit(0)
+    '''
+
+
+
+if __name__ == '__main__':
+   profile_tf_runningmeanstd()