Commit 9f4154e8 authored by Yuxin Wu's avatar Yuxin Wu

allow training tower with index=0 and reuse

parent 17c5984f
...@@ -178,9 +178,10 @@ class TrainTowerContext(BaseTowerContext): ...@@ -178,9 +178,10 @@ class TrainTowerContext(BaseTowerContext):
vs = tf.get_variable_scope() vs = tf.get_variable_scope()
assert vs.name == '', "Cannot nest TrainTowerContext with an existing 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, \ 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 @property
def is_main_training_tower(self): def is_main_training_tower(self):
...@@ -188,6 +189,8 @@ class TrainTowerContext(BaseTowerContext): ...@@ -188,6 +189,8 @@ class TrainTowerContext(BaseTowerContext):
@property @property
def has_own_variables(self): def has_own_variables(self):
if self._original_vs_reuse:
return False
return self.index == 0 or len(self._vs_name) > 0 return self.index == 0 or len(self._vs_name) > 0
def _keys_to_freeze(self): def _keys_to_freeze(self):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment