fix colocation problems

tensorpack/models/regularize.py

@@ -47,12 +47,14 @@ def regularize_cost(regex, func, name='regularize_cost'):
     # If vars are replicated, only regularize those in the current tower
     params = ctx.filter_vars_by_vs_name(params)
 
+    G = tf.get_default_graph()
     with tf.name_scope('regularize_cost'):
         costs = []
         for p in params:
             para_name = p.name
             if re.search(regex, para_name):
-                costs.append(func(p))
+                with G.colocate_with(p):
+                    costs.append(func(p))
                 _log_regularizer(para_name)
         if not costs:
             return tf.constant(0, dtype=tf.float32, name='empty_' + name)

tensorpack/tfutils/gradproc.py

@@ -128,8 +128,7 @@ class MapGradient(GradientProcessor):
         for grad, var in grads:
             if re.match(self.regex, var.op.name):
                 matched = True
-                with tf.device(grad.device):
-                    grad = self.func(grad, var)
+                grad = self.func(grad, var)
                 if grad is not None:
                     ret.append((grad, var))
             else:

tensorpack/tfutils/summary.py

@@ -241,8 +241,9 @@ def add_moving_summary(*args, **kwargs):
                 ema_op = moving_averages.assign_moving_average(
                     ema_var, c, decay,
                     zero_debias=True, name=name + '_EMA_apply')
-            tf.summary.scalar(name + '-summary', ema_op)    # write the EMA value as a summary
             ema_ops.append(ema_op)
+        # cannot add it into colocate group -- will force everything to cpus
+        tf.summary.scalar(name + '-summary', ema_op)    # write the EMA value as a summary
     if coll is not None:
         for op in ema_ops:
             # TODO a new collection to summary every step?