|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from __future__ import absolute_import |
|
from __future__ import division |
|
|
|
import tensorflow as tf |
|
from contextlib import contextmanager |
|
|
|
from tensorflow.python.ops import variable_scope |
|
|
|
|
|
_is_variable_replacing = [False] |
|
|
|
def in_variable_replace_scope(): |
|
return _is_variable_replacing[0] |
|
|
|
@contextmanager |
|
def variable_replace(replacements, no_new=True): |
|
""" A context manager that replaces variables. |
|
|
|
This is a context manager that replaces all calls to |
|
get_variable with the variable in replacements. |
|
This function does not support recursive application. |
|
|
|
Args: |
|
replacements: dict |
|
dictionary mapping a variable to replace (the key), with |
|
the variable one wants to replace this variable with (the value). |
|
no_new: bool |
|
raise an error if variables were created. |
|
This is for sanity checking. |
|
Raises: |
|
ValueError: if a new variable or not all the replacements are used. |
|
""" |
|
|
|
|
|
|
|
|
|
|
|
replacements = {k: v for k, v in replacements.items() if not k == v} |
|
|
|
init_vars = tf.trainable_variables() |
|
old_get_variable = variable_scope.get_variable |
|
old_tf_get_variable = tf.get_variable |
|
|
|
names_replace = {} |
|
has_replaced_names = [] |
|
tf.logging.vlog(2, "Trying to replace") |
|
for k, v in replacements.items(): |
|
tf.logging.vlog(2, k.name + " >> " + v.name) |
|
tf.logging.vlog(2, "===") |
|
|
|
for k, v in replacements.items(): |
|
strip_name = k.name.replace("/read:0", "") |
|
strip_name = strip_name.replace(":0", "") |
|
names_replace[strip_name] = v |
|
|
|
def new_get_variable(name, *args, **kwargs): |
|
|
|
n = tf.get_variable_scope().name + "/" + name |
|
|
|
if n in names_replace: |
|
has_replaced_names.append(n) |
|
return names_replace[n] |
|
else: |
|
return old_get_variable(name, *args, **kwargs) |
|
|
|
|
|
if _is_variable_replacing[0] == True: |
|
raise ValueError("No recursive calling to variable replace allowed.") |
|
|
|
variable_scope.get_variable = new_get_variable |
|
tf.get_variable = new_get_variable |
|
|
|
_is_variable_replacing[0] = True |
|
|
|
yield |
|
|
|
if set(has_replaced_names) != set(names_replace.keys()): |
|
print "Didn't use all replacements" |
|
print "replaced variables that are not requested??" |
|
print "===" |
|
for n in list(set(has_replaced_names) - set(names_replace.keys())): |
|
print n |
|
print "Missed replacing variables" |
|
print "===" |
|
for n in list(set(names_replace.keys()) - set(has_replaced_names)): |
|
print n, "==>", names_replace[n].name |
|
raise ValueError("Fix this -- see stderr") |
|
|
|
|
|
tf.get_variable = old_tf_get_variable |
|
variable_scope.get_variable = old_get_variable |
|
|
|
_is_variable_replacing[0] = False |
|
|
|
final_vars = tf.trainable_variables() |
|
assert set(init_vars) == set(final_vars), "trainable variables changed" |
|
|