allow training tower with index=0 and reuse

tensorpack · May 7, 2020 · 9f4154e · 9f4154e
1 parent 17c5984
commit 9f4154e
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tensorpack/tfutils/tower.py b/tensorpack/tfutils/tower.py
@@ -178,16 +178,19 @@ def __init__(self, ns_name, vs_name='', index=0, total=1):
 
         vs = tf.get_variable_scope()
         assert vs.name == '', "Cannot nest TrainTowerContext with an existing variable scope!"
-        if self.has_own_variables:
+        if vs_name:
             assert not vs.reuse, \
-                "Cannot create tower {} under reuse=True!".format(ns_name)
+                "Cannot create tower {} with vs_name={} under reuse=True!".format(ns_name, vs_name)
+        self._original_vs_reuse = vs.reuse
 
     @property
     def is_main_training_tower(self):
         return self.index == 0
 
     @property
     def has_own_variables(self):
+        if self._original_vs_reuse:
+            return False
         return self.index == 0 or len(self._vs_name) > 0
 
     def _keys_to_freeze(self):