Update the variables in a session
| 40 | |
| 41 | |
| 42 | class SessionUpdate(object): |
| 43 | """ Update the variables in a session """ |
| 44 | |
| 45 | def __init__(self, sess, vars_to_update, ignore_mismatch=False): |
| 46 | """ |
| 47 | Args: |
| 48 | sess (tf.Session): a session object |
| 49 | vars_to_update: a collection of variables to update |
| 50 | ignore_mismatch (bool): ignore failures when the value and the |
| 51 | variable does not match. |
| 52 | """ |
| 53 | self.sess = sess |
| 54 | self.name_map = {v.name: v for v in vars_to_update} |
| 55 | self.ignore_mismatch = ignore_mismatch |
| 56 | |
| 57 | @staticmethod |
| 58 | def relaxed_value_for_var(value, var, ignore_mismatch=False): |
| 59 | """ |
| 60 | Returns a relaxed (possibly reshaped/upcast-ed) version of value, |
| 61 | to be loaded to the given variable. |
| 62 | |
| 63 | Args: |
| 64 | value (ndarray): an numpy array to be loaded to var |
| 65 | var (tf.Variable): |
| 66 | ignore_mismatch (bool): ignore failures when the value and the |
| 67 | variable does not match. |
| 68 | |
| 69 | Returns: |
| 70 | ndarray: a possibly reshaped or casted version of value. |
| 71 | Returns None if `ignore_mismatch==True` and the value and the variable |
| 72 | mismatch. |
| 73 | """ |
| 74 | assert isinstance(var, tf.Variable) |
| 75 | name = var.op.name |
| 76 | |
| 77 | # check incompatible shape |
| 78 | varshape = tuple(var.get_shape().as_list()) |
| 79 | if varshape != value.shape: |
| 80 | if np.prod(varshape) != np.prod(value.shape): |
| 81 | if ignore_mismatch: |
| 82 | logger.warn( |
| 83 | "Cannot load an array of shape {} into variable '{}' whose shape is {}.".format( |
| 84 | value.shape, name, varshape)) |
| 85 | return None |
| 86 | else: |
| 87 | raise ValueError( |
| 88 | "Trying to load an array of shape {} into variable '{}' whose shape is {}.".format( |
| 89 | value.shape, name, varshape)) |
| 90 | # TODO only allow reshape when shape different by empty axis |
| 91 | logger.warn("The tensor is reshaped from {} to {} when assigned to '{}'".format( |
| 92 | value.shape, varshape, name)) |
| 93 | value = value.reshape(varshape) |
| 94 | |
| 95 | # Be permissive, and allow some common type incompatibility problems |
| 96 | def allow_cast(to_type, from_type): |
| 97 | # to_type: a tf dtype |
| 98 | # from_type: a numpy dtype |
| 99 | from_type = tf.as_dtype(from_type) |