def update_metainfo_op_with_vars(metainfo_ph: tf.Tensor, nz_ph: tf.Tensor,
metainfo_var: tf.Variable,
nz_var: tf.Variable) -> tf.Operation:
assign_nz = nz_var.assign(nz_ph)
assign_meta = metainfo_var.assign(metainfo_ph)
with tf.control_dependencies([assign_nz, assign_meta]):
update_op = tf.no_op()
return update_op
def _get_update(self, variable: tf.Variable, gradient: tf.Tensor,
step_size: tf.Tensor) -> List[tf.Operation]:
with tf.variable_scope(variable.op.name):
gradient = tf.cast(gradient, tf.float32)
state_m = tf.get_variable(
"adam_m",
shape=variable.shape,
dtype=variable.dtype,
initializer=tf.zeros_initializer(),
trainable=False,
)
updated_m = (self.beta1 * tf.cast(state_m, tf.float32) +
(1 - self.beta1) * gradient)
state_v = tf.get_variable(
"adam_v",
shape=variable.shape,
dtype=tf.float32,
initializer=tf.zeros_initializer(),
trainable=False,
)
updated_v = self.beta2 * state_v + (1 - self.beta2) * (gradient**2)
delta = step_size * updated_m / (tf.sqrt(updated_v) + self.epsilon)
updated_variable = tf.cast(variable, tf.float32) - delta
return [
variable.assign(tf.cast(updated_variable, variable.dtype)),
state_m.assign(tf.cast(updated_m, state_m.dtype)),
state_v.assign(updated_v),
]
def update_metainfo_op_with_vars(metainfo_ph: tf.Tensor, nz_ph: tf.Tensor,
metainfo_var: tf.Variable,
nz_var: tf.Variable) -> tf.Operation:
"""Returns an op that can be used to update the metainfo on device
:param metainfo_ph: Metainfo placeholder
:param nz_ph: Nonzero-values placeholder
:param metainfo_var: Metainfo variable
:param nz_var: Nonzero-values variable
"""
assign_nz = nz_var.assign(nz_ph)
assign_meta = metainfo_var.assign(metainfo_ph)
with tf.control_dependencies([assign_nz, assign_meta]):
update_op = tf.no_op()
return update_op