def batched_euclidean_distance(y_hat, y, squared=True):
assert y_hat.get_shape().ndims == 3 and y.get_shape().ndims == 3
a = tf.square(tf.reduce_sum(y, axis=2))[:, :, None]
b = tf.square(tf.reduce_sum(y_hat, axis=2))[:, None, :]
D = tf.matmul(y, y_hat, transpose_b=True)
d = a + b - 2 * D
return tf.sqrt(d) if not squared else d
def _observation_cost(obs):
c_theta, s_theta, d_theta = obs[:, :1], obs[:, 1:2], obs[:, 2:3]
theta = tf.math.atan2(s_theta, c_theta)
cost = tf.reduce_sum(tf.square(theta) + 0.1 * tf.square(d_theta),
axis=1)
cost = tf.where(tf.math.is_nan(cost), 1e6 * tf.ones_like(cost), cost)
return cost
def normalized_dist(states):
inner = tf.multiply(states - starting_states, goals - starting_states)
upper = tf.reduce_sum(inner, -1)
sign = tf.sign(upper)
result = sign * tf.square(tf.math.divide(upper, tf.norm(goals - starting_states, ord=2)))
term_1 = tf.square(tf.norm(states - starting_states, 2))
term_2 = tf.square(tf.math.divide(upper, tf.norm(goals - starting_states, ord=2)))
return tf.sqrt(epsilon + tf.abs(result - alpha * (term_1 - term_2)))
def summarize_stats(stats):
"""Summarize a dictionary of variables.
Args:
stats: a dictionary of {name: tensor} to compute stats over.
"""
for name, stat in stats.items():
mean = tf.reduce_mean(stat)
tf.summary.scalar('mean_%s' % name, mean)
tf.summary.scalar('max_%s' % name, tf.reduce_max(stat))
tf.summary.scalar('min_%s' % name, tf.reduce_min(stat))
std = tf.sqrt(tf.reduce_mean(tf.square(stat)) - tf.square(mean) + 1e-10)
tf.summary.scalar('std_%s' % name, std)
tf.summary.histogram(name, stat)
def train_step(self,
time_step: ActionTimeStep,
state,
calc_intrinsic_reward=True):
"""
Args:
time_step (ActionTimeStep): input time_step data for ICM
state (Tensor): state for ICM (previous observation)
calc_intrinsic_reward (bool): if False, only return the losses
Returns:
TrainStep:
outputs: empty tuple ()
state: observation
info (ICMInfo):
"""
feature = time_step.observation
prev_action = time_step.prev_action
if self._encoding_net is not None:
feature, _ = self._encoding_net(feature)
prev_feature = state
prev_action = self._encode_action(prev_action)
forward_pred, _ = self._forward_net(
inputs=[tf.stop_gradient(prev_feature), prev_action])
forward_loss = 0.5 * tf.reduce_mean(
tf.square(tf.stop_gradient(feature) - forward_pred), axis=-1)
action_pred, _ = self._inverse_net(inputs=[prev_feature, feature])
if tensor_spec.is_discrete(self._action_spec):
inverse_loss = tf.nn.softmax_cross_entropy_with_logits(
labels=prev_action, logits=action_pred)
else:
inverse_loss = 0.5 * tf.reduce_mean(
tf.square(prev_action - action_pred), axis=-1)
intrinsic_reward = ()
if calc_intrinsic_reward:
intrinsic_reward = tf.stop_gradient(forward_loss)
intrinsic_reward = self._reward_normalizer.normalize(
intrinsic_reward)
return AlgorithmStep(
outputs=(),
state=feature,
info=ICMInfo(reward=intrinsic_reward,
loss=LossInfo(loss=forward_loss + inverse_loss,
extra=dict(forward_loss=forward_loss,
inverse_loss=inverse_loss))))
def gauss_kernel2D(x, Dx, Dy, gamma=1.):
h_size = (x.get_shape()[-1].value) // 2
x = tf.expand_dims(x, axis=-1)
if x.get_shape().ndims < 4:
Dx = tf.reshape(Dx, (1, 1, -1))
Dy = tf.reshape(Dy, (1, 1, -1))
x1, x2 = x[:, :h_size], x[:, h_size:]
else:
Dy = tf.reshape(Dy, (1, 1, 1, 1, -1))
Dx = tf.reshape(Dx, (1, 1, 1, 1, -1))
x1, x2 = x[:, :, :, :h_size], x[:, :, :, h_size:]
gauss_kernel = tf.exp(-gamma * tf.square(x1 - Dx)) + tf.exp(- gamma * tf.square(x2 - Dy))
return gauss_kernel
def compute_logits(self,
support_embeddings,
query_embeddings,
onehot_support_labels,
cosine_distance=False):
"""Computes the negative distances of each query point to each prototype."""
prototypes = compute_prototypes(support_embeddings,
onehot_support_labels)
if cosine_distance:
query_embeddings = tf.nn.l2_normalize(query_embeddings,
1,
epsilon=1e-3)
prototypes = tf.nn.l2_normalize(prototypes, 1, epsilon=1e-3)
logits = tf.matmul(query_embeddings, prototypes, transpose_b=True)
else:
# [num test images, 1, embedding size].
query_embeddings = tf.expand_dims(query_embeddings, 1)
# [1, num_clases, embedding_size].
prototypes = tf.expand_dims(prototypes, 0)
# Squared euclidean distance between each test embedding / prototype pair.
distances = tf.reduce_sum(tf.square(query_embeddings - prototypes),
2)
logits = -distances
return logits
def train_step(self, inputs, state, calc_intrinsic_reward=True):
"""
Args:
inputs (tuple): observation
state (tuple): empty tuple ()
calc_intrinsic_reward (bool): if False, only return the losses
Returns:
TrainStep:
outputs: empty tuple ()
state: empty tuple ()
info: RNDInfo
"""
observation, _ = inputs
if self._observation_normalizer is not None:
observation = self._observation_normalizer.normalize(observation)
pred_embedding, _ = self._predictor_net(observation)
target_embedding, _ = self._target_net(observation)
loss = 0.5 * tf.reduce_mean(
tf.square(pred_embedding - tf.stop_gradient(target_embedding)),
axis=-1)
intrinsic_reward = ()
if calc_intrinsic_reward:
intrinsic_reward = tf.stop_gradient(loss)
intrinsic_reward = self._reward_normalizer.normalize(
intrinsic_reward)
return AlgorithmStep(outputs=(),
state=(),
info=RNDInfo(reward=intrinsic_reward,
loss=LossInfo(loss=loss)))
def get_train_op():
loss = tf.reduce_mean(input_tensor=tf.square(q_clicked - target_clicked))
if self.summary_writer is not None:
with tf.variable_scope('Losses'):
tf.summary.scalar('Loss', loss)
return loss
def ctrl_rewards(states,
actions,
rewards,
next_states,
contexts,
reward_scales=1.0):
"""Returns the negative control cost.
Args:
states: A [batch_size, num_state_dims] Tensor representing a batch
of states.
actions: A [batch_size, num_action_dims] Tensor representing a batch
of actions.
rewards: A [batch_size] Tensor representing a batch of rewards.
next_states: A [batch_size, num_state_dims] Tensor representing a batch
of next states.
contexts: A list of [batch_size, num_context_dims] Tensor representing
a batch of contexts.
reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
must be broadcastable to number of reward dimensions.
Returns:
A new tf.float32 [batch_size] rewards Tensor, and
tf.float32 [batch_size] discounts tensor.
"""
del states, rewards, contexts # Unused
if actions is None:
rewards = tf.to_float(tf.zeros(shape=next_states.shape[:1]))
else:
rewards = -tf.reduce_sum(tf.square(actions), axis=1)
rewards *= reward_scales
rewards = tf.to_float(rewards)
return rewards, tf.ones_like(rewards)
def __init__(self, input_shape, dict_size=(-1., 1., 20), gamma=None):
self.d = tf.linspace(*dict_size)
if gamma is None:
self.gamma = .5 / tf.square(2 * (self.d[-1] - self.d[0])) # (d_stop - d_start)*2
else:
self.gamma = gamma
self.alpha = tf.get_variable('alpha', shape=(1, input_shape, self.d.get_shape()[0]),
initializer=RidgeInit(gauss_kernel(self.d, self.d, self.gamma), self.d))
def gauss_kernel(x, D, gamma=1.):
x = tf.expand_dims(x, axis=-1)
if x.get_shape().ndims < 4:
D = tf.reshape(D, (1, 1, -1))
else:
D = tf.reshape(D, (1, 1, 1, 1, -1))
return tf.exp(- gamma * tf.square(x - D))
def _reg(cls, batch_size, d, x, x_fake, beta=1e-1):
alpha = tf.random_uniform(shape=[batch_size, 1], minval=0., maxval=1.)
interpolates = alpha * x + (1 - alpha) * x_fake
int_d = d(interpolates)
gradients = tf.gradients(int_d, [interpolates])[0]
slopes = tf.sqrt(
tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
return beta * tf.reduce_mean((slopes - 1)**2)
def _variable_summaries(var):
"""Attach a lot of summaries to a Tensor (for TensorBoard visualization)."""
with tf.name_scope('summaries'):
mean = tf.reduce_mean(var)
tf.summary.scalar('mean', mean)
with tf.name_scope('stddev'):
stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
tf.summary.scalar('stddev', stddev)
tf.summary.scalar('max', tf.reduce_max(var))
tf.summary.scalar('min', tf.reduce_min(var))
tf.summary.histogram('histogram', var)
def train_step(self, inputs, state):
"""
Args:
inputs (tuple): observation and previous action
Returns:
TrainStep:
outputs: intrinsic reward
state:
info:
"""
feature, prev_action = inputs
if self._encoding_net is not None:
feature, _ = self._encoding_net(feature)
prev_feature = state
prev_action = self._encode_action(prev_action)
forward_pred, _ = self._forward_net(
inputs=[tf.stop_gradient(prev_feature), prev_action])
forward_loss = 0.5 * tf.reduce_mean(
tf.square(tf.stop_gradient(feature) - forward_pred), axis=-1)
action_pred, _ = self._inverse_net(inputs=[prev_feature, feature])
if tensor_spec.is_discrete(self._action_spec):
inverse_loss = tf.nn.softmax_cross_entropy_with_logits(
labels=prev_action, logits=action_pred)
else:
inverse_loss = 0.5 * tf.reduce_mean(
tf.square(prev_action - action_pred), axis=-1)
intrinsic_reward = tf.stop_gradient(forward_loss)
intrinsic_reward = self._reward_normalizer.normalize(intrinsic_reward)
return AlgorithmStep(outputs=intrinsic_reward,
state=feature,
info=LossInfo(loss=forward_loss + inverse_loss,
extra=ICMLossInfo(
forward_loss=forward_loss,
inverse_loss=inverse_loss)))
def compute_logits(self, support_embeddings, query_embeddings,
onehot_support_labels):
"""Computes the negative distances of each query point to each prototype."""
# [num test images, 1, embedding size].
query_embeddings = tf.expand_dims(query_embeddings, 1)
prototypes = compute_prototypes(support_embeddings, onehot_support_labels)
# [1, num_clases, embedding_size].
prototypes = tf.expand_dims(prototypes, 0)
# Squared euclidean distances between each test embedding / prototype pair.
distances = tf.reduce_sum(tf.square(query_embeddings - prototypes), 2)
return -distances
def proto_maml_fc_bias(self, prototypes, zero_pad_to_max_way=False):
"""Computes the Prototypical MAML fc layer's bias.
Args:
prototypes: Tensor of shape [num_classes, embedding_size]
zero_pad_to_max_way: Whether to zero padd to max num way.
Returns:
fc_bias: Tensor of shape [num_classes] or [self.logit_dim]
when zero_pad_to_max_way is True.
"""
fc_bias = -tf.square(tf.norm(prototypes, axis=1))
if zero_pad_to_max_way:
paddings = [[0, self.logit_dim - tf.shape(fc_bias)[0]]]
fc_bias = tf.pad(fc_bias, paddings, 'CONSTANT', constant_values=0)
return fc_bias
def train_step(self,
time_step: ActionTimeStep,
state,
calc_intrinsic_reward=True):
"""
Args:
time_step (ActionTimeStep): input time_step data
state (tuple): empty tuple ()
calc_intrinsic_reward (bool): if False, only return the losses
Returns:
TrainStep:
outputs: empty tuple ()
state: empty tuple ()
info: ICMInfo
"""
observation = time_step.observation
if self._stacked_frames:
# Assuming stacking in the last dim, we only keep the last frame.
observation = observation[..., -1:]
if self._observation_normalizer is not None:
observation = self._observation_normalizer.normalize(observation)
if self._encoder_net is not None:
observation = tf.stop_gradient(self._encoder_net(observation)[0])
pred_embedding, _ = self._predictor_net(observation)
target_embedding, _ = self._target_net(observation)
loss = tf.reduce_sum(
tf.square(pred_embedding - tf.stop_gradient(target_embedding)),
axis=-1)
intrinsic_reward = ()
if calc_intrinsic_reward:
intrinsic_reward = tf.stop_gradient(loss)
if self._reward_normalizer:
intrinsic_reward = self._reward_normalizer.normalize(
intrinsic_reward, clip_value=self._reward_clip_value)
return AlgorithmStep(
outputs=(),
state=(),
info=ICMInfo(reward=intrinsic_reward, loss=LossInfo(loss=loss)))
def _get_dist(self, query_queries, query_values, support_keys, support_values,
labels):
"""Get distances between queries and query-aligned prototypes."""
# attended_values: [N_support, n_query, h_query, w_query, C]
attended_values = self._attend(query_queries, support_keys, support_values,
labels)
# query_aligned_prototypes: [N_classes, n_query, h_query, w_query, C]
query_aligned_prototypes = tf.math.unsorted_segment_sum(
attended_values, labels,
tf.reduce_max(labels) + 1)
# (scaled) Euclidean distance
shp = tf.shape(query_values)
aligned_dist = tf.square(query_values[tf.newaxis, Ellipsis] -
query_aligned_prototypes)
return tf.reduce_sum(aligned_dist, [2, 3, 4]) / tf.cast(
shp[-3] * shp[-2], aligned_dist.dtype)
def train_step(self, time_step: ActionTimeStep, state: DynamicsState):
"""
Args:
time_step (ActionTimeStep): input data for dynamics learning
state (Tensor): state for dynamics learning (previous observation)
Returns:
TrainStep:
outputs: empty tuple ()
state (DynamicsState): state for training
info (DynamicsInfo):
"""
feature = time_step.observation
dynamics_step = self.predict(time_step, state)
forward_pred = dynamics_step.outputs
forward_loss = 0.5 * tf.reduce_mean(
tf.square(feature - forward_pred), axis=-1)
info = DynamicsInfo(
loss=LossInfo(
loss=forward_loss, extra=dict(forward_loss=forward_loss)))
state = DynamicsState(feature=feature)
return AlgorithmStep(outputs=(), state=state, info=info)
def l2_norm(x, axis=2):
squared = tf.reduce_sum(tf.square(x), axis=axis, keepdims=True)
norm = tf.sqrt(tf.maximum(squared, 1e-6))
return norm
def projection_distance(states,
starting_states,
actions,
rewards,
next_states,
contexts,
alpha = 0,
state_scales=1.0,
goal_scales=1.0,
reward_scales=1.0,
weight_index=None,
weight_vector=None,
summarize=False,
termination_epsilon=1e-4,
state_indices=None,
goal_indices=None,
vectorize=False,
relative_context=False,
diff=False,
norm='L2',
epsilon=1e-10,
bonus_epsilon=0., #5.,
offset=0.0):
"""Returns the negative euclidean distance between next_states and contexts.
Args:
states: A [batch_size, num_state_dims] Tensor representing a batch
of states.
actions: A [batch_size, num_action_dims] Tensor representing a batch
of actions.
rewards: A [batch_size] Tensor representing a batch of rewards.
next_states: A [batch_size, num_state_dims] Tensor representing a batch
of next states.
contexts: A list of [batch_size, num_context_dims] Tensor representing
a batch of contexts.
state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
must be broadcastable to number of state dimensions.
goal_scales: multiplicative scale for goals. A scalar or 1D tensor,
must be broadcastable to number of goal dimensions.
reward_scales: multiplicative scale for rewards. A scalar or 1D tensor,
must be broadcastable to number of reward dimensions.
weight_index: (integer) The context list index that specifies weight.
weight_vector: (a number or a list or Numpy array) The weighting vector,
broadcastable to `next_states`.
summarize: (boolean) enable summary ops.
termination_epsilon: terminate if dist is less than this quantity.
state_indices: (a list of integers) list of state indices to select.
goal_indices: (a list of integers) list of goal indices to select.
vectorize: Return a vectorized form.
norm: L1 or L2.
epsilon: small offset to ensure non-negative/zero distance.
Returns:
A new tf.float32 [batch_size] rewards Tensor, and
tf.float32 [batch_size] discounts tensor.
"""
del actions, rewards # Unused
stats = {}
record_tensor(next_states, state_indices, stats, 'next_states')
states = index_states(states, state_indices)
starting_states = index_states(starting_states, state_indices)
next_states = index_states(next_states, state_indices)
goals = index_states(contexts[0], goal_indices)
if relative_context:
goals = states + goals
sq_dists = tf.squared_difference(next_states * state_scales,
goals * goal_scales)
dist = tf.reduce_sum(sq_dists, -1)
#def normalized_dist(states):
# dot_product = tf.matmul(states - starting_states, tf.transpose(goals - starting_states))
# return goals - starting_states - dot_product
def projection_dist(states):
inner = tf.multiply(states - starting_states, goals - starting_states)
upper = tf.reduce_sum(inner, -1)
sign = tf.sign(upper)
result = tf.math.divide(upper, tf.norm(goals - starting_states, ord=2))
term_1 = tf.norm(states - starting_states, 2)
return -1*term_1+result
dist_s = projection_dist(states)
dist_s = tf.sqrt(tf.square(dist_s) + epsilon)
dist_ns = projection_dist(next_states)
ret = dist_ns, tf.to_float(dist > termination_epsilon)
return ret
def _squared_frobenius_norm(t):
return tf.reduce_sum(tf.square(t))
def _action_cost(action):
return 0.001 * tf.reduce_sum(tf.square(action), axis=1)
def bn(x,
params=None,
moments=None,
backprop_through_moments=True,
use_ema=False,
is_training=True,
ema_epsilon=.9):
"""Batch normalization.
The usage should be as follows: If x is the support images, moments should be
None so that they are computed from the support set examples. On the other
hand, if x is the query images, the moments argument should be used in order
to pass in the mean and var that were computed from the support set.
Args:
x: inputs.
params: None or a dict containing the values of the offset and scale params.
moments: None or a dict containing the values of the mean and var to use for
batch normalization.
backprop_through_moments: Whether to allow gradients to flow through the
given support set moments. Only applies to non-transductive batch norm.
use_ema: apply moving averages of batch norm statistics, or update them,
depending on whether we are training or testing. Note that passing
moments will override this setting, and result in neither updating or
using ema statistics. This is important to make sure that episodic
learners don't update ema statistics a second time when processing
queries.
is_training: if use_ema=True, this determines whether to apply the moving
averages, or update them.
ema_epsilon: if updating moving averages, use this value for the
exponential moving averages.
Returns:
output: The result of applying batch normalization to the input.
params: The updated params.
moments: The updated moments.
"""
params_keys, params_vars, moments_keys, moments_vars = [], [], [], []
with tf.variable_scope('batch_norm'):
scope_name = tf.get_variable_scope().name
if use_ema:
ema_shape = [1, 1, 1, x.get_shape().as_list()[-1]]
mean_ema = tf.get_variable(
'mean_ema',
shape=ema_shape,
initializer=tf.initializers.zeros(),
trainable=False)
var_ema = tf.get_variable(
'var_ema',
shape=ema_shape,
initializer=tf.initializers.ones(),
trainable=False)
if moments is not None:
if backprop_through_moments:
mean = moments[scope_name + '/mean']
var = moments[scope_name + '/var']
else:
# This variant does not yield good resutls.
mean = tf.stop_gradient(moments[scope_name + '/mean'])
var = tf.stop_gradient(moments[scope_name + '/var'])
elif use_ema and not is_training:
mean = mean_ema
var = var_ema
else:
# If not provided, compute the mean and var of the current batch.
replica_ctx = tf.distribute.get_replica_context()
if replica_ctx:
# from third_party/tensorflow/python/keras/layers/normalization_v2.py
axes = list(range(len(x.shape) - 1))
local_sum = tf.reduce_sum(x, axis=axes, keepdims=True)
local_squared_sum = tf.reduce_sum(
tf.square(x), axis=axes, keepdims=True)
batch_size = tf.cast(tf.shape(x)[0], tf.float32)
x_sum, x_squared_sum, global_batch_size = (
replica_ctx.all_reduce('sum',
[local_sum, local_squared_sum, batch_size]))
axes_vals = [(tf.shape(x))[i] for i in range(1, len(axes))]
multiplier = tf.cast(tf.reduce_prod(axes_vals), tf.float32)
multiplier = multiplier * global_batch_size
mean = x_sum / multiplier
x_squared_mean = x_squared_sum / multiplier
# var = E(x^2) - E(x)^2
var = x_squared_mean - tf.square(mean)
else:
mean, var = tf.nn.moments(
x, axes=list(range(len(x.shape) - 1)), keep_dims=True)
# Only update ema's if training and we computed the moments in the current
# call. Note: at test time for episodic learners, ema's may be passed
# from the support set to the query set, even if it's not really needed.
if use_ema and is_training and moments is None:
replica_ctx = tf.distribute.get_replica_context()
mean_upd = tf.assign(mean_ema,
mean_ema * ema_epsilon + mean * (1.0 - ema_epsilon))
var_upd = tf.assign(var_ema,
var_ema * ema_epsilon + var * (1.0 - ema_epsilon))
updates = tf.group([mean_upd, var_upd])
if replica_ctx:
tf.add_to_collection(
tf.GraphKeys.UPDATE_OPS,
tf.cond(
tf.equal(replica_ctx.replica_id_in_sync_group, 0),
lambda: updates, tf.no_op))
else:
tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, updates)
moments_keys += [scope_name + '/mean']
moments_vars += [mean]
moments_keys += [scope_name + '/var']
moments_vars += [var]
if params is None:
offset = tf.get_variable(
'offset',
shape=mean.get_shape().as_list(),
initializer=tf.initializers.zeros())
scale = tf.get_variable(
'scale',
shape=var.get_shape().as_list(),
initializer=tf.initializers.ones())
else:
offset = params[scope_name + '/offset']
scale = params[scope_name + '/scale']
params_keys += [scope_name + '/offset']
params_vars += [offset]
params_keys += [scope_name + '/scale']
params_vars += [scale]
output = tf.nn.batch_normalization(x, mean, var, offset, scale, 0.00001)
params = collections.OrderedDict(zip(params_keys, params_vars))
moments = collections.OrderedDict(zip(moments_keys, moments_vars))
return output, params, moments
def linear_classifier(embeddings, num_classes, cosine_classifier,
cosine_logits_multiplier, use_weight_norm, weight_decay):
"""Forward pass through a linear classifier, or possibly a cosine classifier.
Args:
embeddings: A Tensor of size [batch size, embedding dim].
num_classes: An integer; the dimension of the classification.
cosine_classifier: A bool. If true, a cosine classifier is used, which does
not require a bias.
cosine_logits_multiplier: A float. Only used if cosine_classifier is True,
and multiplies the resulting logits.
use_weight_norm: A bool. Whether weight norm was used. If so, then if using
cosine classifier, normalize only the embeddings but not the weights.
weight_decay: A float; the scalar multiple on the L2 regularization of the
weight matrix.
Returns:
logits: A Tensor of size [batch size, num outputs].
"""
embedding_dims = embeddings.get_shape().as_list()[-1]
if use_weight_norm:
# A variable to keep track of whether the initialization has already
# happened.
data_dependent_init_done = tf.get_variable('data_dependent_init_done',
initializer=0,
dtype=tf.int32,
trainable=False)
w_fc = tf.get_variable('w_fc', [embedding_dims, num_classes],
initializer=tf.random_normal_initializer(
0, 0.05),
trainable=True)
# This init is temporary as it needs to be done in a data-dependent way.
# It will be overwritten during the first forward pass through this layer.
g = tf.get_variable('g',
dtype=tf.float32,
initializer=tf.ones([num_classes]),
trainable=True)
b_fc = None
if not cosine_classifier:
# Also initialize a bias.
b_fc = tf.get_variable('b_fc',
initializer=tf.zeros([num_classes]),
trainable=True)
def _do_data_dependent_init():
"""Returns ops for the data-dependent init of g and maybe b_fc."""
w_fc_normalized = tf.nn.l2_normalize(w_fc.read_value(), [0])
output_init = tf.matmul(embeddings, w_fc_normalized)
mean_init, var_init = tf.nn.moments(output_init, [0])
# Data-dependent init values.
g_init_value = 1. / tf.sqrt(var_init + 1e-10)
ops = [tf.assign(g, g_init_value)]
if not cosine_classifier:
# Also initialize a bias in a data-dependent way.
b_fc_init_value = -mean_init * g_init_value
ops.append(tf.assign(b_fc, b_fc_init_value))
# Mark that the data-dependent initialization is done to prevent it from
# happening again in the future.
ops.append(tf.assign(data_dependent_init_done, 1))
return tf.group(*ops)
# Possibly perform data-dependent init (if it hasn't been done already).
init_op = tf.cond(tf.equal(data_dependent_init_done, 0),
_do_data_dependent_init, tf.no_op)
with tf.control_dependencies([init_op]):
# Apply weight normalization.
w_fc *= g / tf.sqrt(tf.reduce_sum(tf.square(w_fc), [0]))
# Forward pass through the layer defined by w_fc and b_fc.
logits = linear_classifier_forward_pass(embeddings, w_fc, b_fc,
cosine_classifier,
cosine_logits_multiplier,
True)
else:
# No weight norm.
w_fc = functional_backbones.weight_variable(
[embedding_dims, num_classes], weight_decay=weight_decay)
b_fc = None
if not cosine_classifier:
# Also initialize a bias.
b_fc = functional_backbones.bias_variable([num_classes])
# Forward pass through the layer defined by w_fc and b_fc.
logits = linear_classifier_forward_pass(embeddings, w_fc, b_fc,
cosine_classifier,
cosine_logits_multiplier,
False)
return logits
def stability_loss(h, beta):
if beta == 0.0:
return 0.0
else:
l2 = tf.sqrt(tf.reduce_sum(tf.square(h), axis=-1))
return beta * tf.reduce_mean(tf.square(l2[1:] - l2[:-1]))
def euclidean_distance(x, y): """Computes the Euclidean distance.""" x = tf.expand_dims(x, 1) y = tf.expand_dims(y, 0) return tf.reduce_sum(input_tensor=tf.square(x - y), axis=2)
def mse(y, y_hat):
return tf.square(y - y_hat, name="mse")
