def actor_loss_fn(dqda, action):
if self._dqda_clipping:
dqda = tf.clip_by_value(dqda, -self._dqda_clipping,
self._dqda_clipping)
loss = 0.5 * losses.element_wise_squared_loss(
tf.stop_gradient(dqda + action), action)
loss = tf.reduce_sum(loss, axis=list(range(1, len(loss.shape))))
return loss
def __call__(self, training_info: TrainingInfo, value):
"""Cacluate actor critic loss
The first dimension of all the tensors is time dimension and the second
dimesion is the batch dimension.
Args:
training_info (TrainingInfo): training_info collected by
(On/Off)PolicyDriver. All tensors in training_info are time-major
value (tf.Tensor): the time-major tensor for the value at each time
step
final_value (tf.Tensor): the value at one step ahead.
Returns:
loss_info (LossInfo): with loss_info.extra being ActorCriticLossInfo
"""
returns, advantages = self._calc_returns_and_advantages(
training_info, value)
def _summary():
with tf.name_scope('ActorCriticLoss'):
tf.summary.scalar("values", tf.reduce_mean(value))
tf.summary.scalar("returns", tf.reduce_mean(returns))
tf.summary.scalar("advantages/mean",
tf.reduce_mean(advantages))
tf.summary.histogram("advantages/value", advantages)
tf.summary.scalar("explained_variance_of_return_by_value",
common.explained_variance(value, returns))
if self._debug_summaries:
common.run_if(common.should_record_summaries(), _summary)
if self._normalize_advantages:
advantages = _normalize_advantages(advantages, axes=(0, 1))
if self._advantage_clip:
advantages = tf.clip_by_value(advantages, -self._advantage_clip,
self._advantage_clip)
pg_loss = self._pg_loss(training_info, tf.stop_gradient(advantages))
td_loss = self._td_error_loss_fn(tf.stop_gradient(returns), value)
loss = pg_loss + self._td_loss_weight * td_loss
entropy_loss = ()
if self._entropy_regularization is not None:
entropy, entropy_for_gradient = dist_utils.entropy_with_fallback(
training_info.action_distribution, self._action_spec)
entropy_loss = -entropy
loss -= self._entropy_regularization * entropy_for_gradient
return LossInfo(loss=loss,
extra=ActorCriticLossInfo(td_loss=td_loss,
pg_loss=pg_loss,
entropy_loss=entropy_loss))
def loss_fn(self):
adv = tf.placeholder(tf.float32, [None], name="advantages")
returns = tf.placeholder(tf.float32, [None], name="returns")
logli_old = tf.placeholder(tf.float32, [None], name="logli_old")
value_old = tf.placeholder(tf.float32, [None], name="value_old")
ratio = tf.exp(self.policy.logli - logli_old)
clipped_ratio = tf.clip_by_value(ratio, 1-self.clip_ratio, 1+self.clip_ratio)
value_err = (self.value - returns)**2
if self.clip_value > 0.0:
clipped_value = tf.clip_by_value(self.value, value_old-self.clip_value, value_old+self.clip_value)
clipped_value_err = (clipped_value - returns)**2
value_err = tf.maximum(value_err, clipped_value_err)
policy_loss = -tf.reduce_mean(tf.minimum(adv * ratio, adv * clipped_ratio))
value_loss = tf.reduce_mean(value_err) * self.value_coef
entropy_loss = tf.reduce_mean(self.policy.entropy) * self.entropy_coef
# we want to reduce policy and value errors, and maximize entropy
# but since optimizer is minimizing the signs are opposite
full_loss = policy_loss + value_loss - entropy_loss
return full_loss, [policy_loss, value_loss, entropy_loss], [adv, returns, logli_old, value_old]
def call(self, observation, step_type=(), network_state=()):
del step_type # unused.
output = tf.cast(tf.nest.flatten(observation)[0], tf.float32)
for layer in self._mlp_layers:
output = layer(output)
shift, log_scale_diag = tf.split(output, 2, axis=-1)
log_scale_diag = tf.clip_by_value(log_scale_diag, -20, 2)
base_distribution = tfp.distributions.MultivariateNormalDiag(
loc=shift, scale_diag=tf.exp(log_scale_diag))
distribution = SquashToSpecDistribution(
base_distribution, self._single_action_spec)
distribution = tf.nest.pack_sequence_as(self.output_spec, [distribution])
return distribution, network_state
def loss_fn(self):
adv = tf.placeholder(tf.float32, [None], name="advantages")
returns = tf.placeholder(tf.float32, [None], name="returns")
logli_old = tf.placeholder(tf.float32, [None], name="logli_old")
ratio = tf.exp(self.policy.logli - logli_old)
clipped_ratio = tf.clip_by_value(ratio, 1-self.clip_ratio, 1+self.clip_ratio)
policy_loss = -tf.reduce_mean(tf.minimum(adv * ratio, adv * clipped_ratio))
# TODO clip value loss
value_loss = tf.reduce_mean((self.value - returns)**2) * self.value_coef
entropy_loss = tf.reduce_mean(self.policy.entropy) * self.entropy_coef
# we want to reduce policy and value errors, and maximize entropy
# but since optimizer is minimizing the signs are opposite
full_loss = policy_loss + value_loss - entropy_loss
return full_loss, [policy_loss, value_loss, entropy_loss], [adv, returns, logli_old]
def optimizer_update(iterate_collection, iteration_idx, objective_fn,
update_fn, get_params_fn, first_order, clip_grad_norm):
"""Returns the next iterate in the optimization of objective_fn wrt variables.
Args:
iterate_collection: A (potentially structured) container of tf.Tensors
corresponding to the state of the current iterate.
iteration_idx: An int Tensor; the iteration number.
objective_fn: Callable that takes in variables and produces the value of the
objective function.
update_fn: Callable that takes in the gradient of the objective function and
the current iterate and produces the next iterate.
get_params_fn: Callable that takes in the gradient of the objective function
and the current iterate and produces the next iterate.
first_order: If True, prevent the computation of higher order gradients.
clip_grad_norm: If not None, gradient dimensions are independently clipped
to lie in the interval [-clip_grad_norm, clip_grad_norm].
"""
variables = [get_params_fn(iterate) for iterate in iterate_collection]
if tf.executing_eagerly():
with tf.GradientTape(persistent=True) as g:
g.watch(variables)
loss = objective_fn(variables, iteration_idx)
grads = g.gradient(loss, variables)
else:
loss = objective_fn(variables, iteration_idx)
grads = tf.gradients(ys=loss, xs=variables)
if clip_grad_norm:
grads = [
tf.clip_by_value(grad, -1 * clip_grad_norm, clip_grad_norm)
for grad in grads
]
if first_order:
grads = [tf.stop_gradient(dv) for dv in grads]
return [
update_fn(i=iteration_idx, grad=dv, state=s)
for (s, dv) in zip(iterate_collection, grads)
]
def train_step(self,
time_step: ActionTimeStep,
state,
calc_intrinsic_reward=True):
"""
Args:
time_step (ActionTimeStep): input time_step data
state (tuple): state for MISC (previous observation,
previous previous action)
calc_intrinsic_reward (bool): if False, only return the losses
Returns:
TrainStep:
outputs: empty tuple ()
state: tuple of observation and previous action
info: (MISCInfo):
"""
feature_state = time_step.observation
prev_action = time_step.prev_action
feature = tf.concat([feature_state, prev_action], axis=-1)
prev_feature = tf.concat(state, axis=-1)
feature_reshaped = tf.expand_dims(feature, axis=1)
prev_feature_reshaped = tf.expand_dims(prev_feature, axis=1)
feature_pair = tf.concat([prev_feature_reshaped, feature_reshaped], 1)
feature_reshaped_tran = transpose2(feature_reshaped, 1, 0)
def add_batch():
self._buffer.add_batch(feature_reshaped_tran)
# The reason that the batch_size of time_step can be different from
# self._batch_size is that this is invoked in PREPARE_SPEC mode.
# TODO: handle PREPARE_SPEC properly
if (calc_intrinsic_reward
and time_step.step_type.shape[0] == self._batch_size):
add_batch()
if self._n_objects < 2:
obs_tau_excludes_goal, obs_tau_achieved_goal = (
self._split_observation_fn(feature_pair))
loss = self._mine(obs_tau_excludes_goal, obs_tau_achieved_goal)
elif self._n_objects == 2:
(obs_tau_excludes_goal, obs_tau_achieved_goal_1,
obs_tau_achieved_goal_2
) = self._split_observation_fn(feature_pair)
loss_1 = self._mine(obs_tau_excludes_goal, obs_tau_achieved_goal_1)
loss_2 = self._mine(obs_tau_excludes_goal, obs_tau_achieved_goal_2)
loss = loss_1 + loss_2
intrinsic_reward = ()
if calc_intrinsic_reward:
# scale/normalize the MISC intrinsic reward
if self._n_objects < 2:
intrinsic_reward = tf.clip_by_value(self._mi_r_scale * loss, 0,
1)
elif self._n_objects == 2:
intrinsic_reward = tf.clip_by_value(
self._mi_r_scale * loss_1, 0,
1) + 1 * tf.clip_by_value(self._mi_r_scale * loss_2, 0, 1)
return AlgorithmStep(outputs=(),
state=[feature_state, prev_action],
info=MISCInfo(reward=intrinsic_reward))
def photometric_augmentation(
images,
probability_color_swap = 0.0,
probability_hue_shift = 1.0,
probability_saturation = 1.0,
probability_brightness = 1.0,
probability_contrast = 1.0,
probability_gaussian_noise = 0.0,
probability_brightness_individual = 0.0,
probability_contrast_individual = 0.0,
probability_eraser = 0.5,
probability_eraser_additional_operations = 0.5,
probability_assymetric = 0.2,
max_delta_hue = 0.5 / 3.14,
min_bound_saturation = 0.6,
max_bound_saturation = 1.4,
max_delta_brightness = 0.4,
min_bound_contrast = 0.6,
max_bound_contrast = 1.4,
min_bound_gaussian_noise = 0.0,
max_bound_gaussian_noise = 0.02,
max_delta_brightness_individual = 0.02,
min_bound_contrast_individual = 0.95,
max_bound_contrast_individual = 1.05,
min_size_eraser = 50,
max_size_eraser = 100,
max_operations_eraser = 3):
"""Applies photometric augmentations to an image pair.
Args:
images: Image pair of shape [2, height, width, channels].
probability_color_swap: Probability of applying color swap augmentation.
probability_hue_shift: Probability of applying hue shift augmentation.
probability_saturation: Probability of applying saturation augmentation.
probability_brightness: Probability of applying brightness augmentation.
probability_contrast: Probability of applying contrast augmentation.
probability_gaussian_noise: Probability of applying gaussian noise
augmentation.
probability_brightness_individual: Probability of applying brightness
augmentation individually to each image of the image pair.
probability_contrast_individual: Probability of applying contrast
augmentation individually to each image of the image pair.
probability_eraser: Probability of applying the eraser augmentation.
probability_eraser_additional_operations: Probability of applying additional
erase operations within the eraser augmentation.
probability_assymetric: Probability of applying some photomoteric
augmentations invidiually per frame (hue_shift, brightness,
saturation, contrast, gaussian noise).
max_delta_hue: Must be in the interval [0, 0.5]. Defines the interval
[-max_delta_hue, max_delta_hue] in which to pick a random hue offset.
min_bound_saturation: Lower bound for the randomly picked saturation factor
in the interval [lower, upper].
max_bound_saturation: Upper bound for the randomly picked saturation factor
in the interval [lower, upper].
max_delta_brightness: Delta defines the interval [-max_delta, max_delta) in
which a random value is picked that will be added to the image values.
min_bound_contrast: Lower bound for the randomly picked contrast factor in
the interval [lower, upper]. It will be applied per channel via (x - mean)
* contrast_factor + mean.
max_bound_contrast: Upper bound for the randomly picked contrast factor in
the interval [lower, upper]. It will be applied per channel via (x - mean)
* contrast_factor + mean.
min_bound_gaussian_noise: Lower bound for the randomly picked sigma in the
interval [lower, upper].
max_bound_gaussian_noise: Upper bound for the randomly picked sigma in the
interval [lower, upper].
max_delta_brightness_individual: Same as max_delta_brightness, but for the
augmentation applied individually per frame.
min_bound_contrast_individual: Same as min_bound_contrast, but for the
augmentation applied individually per frame.
max_bound_contrast_individual: Same as max_bound_contrast, but for the
augmentation applied individually per frame.
min_size_eraser: Minimal side length of the rectangle shaped region that
will be removed.
max_size_eraser: Maximal side length of the rectangle shaped region that
will be removed.
max_operations_eraser: Maximal number of rectangle shaped regions that will
be removed.
Returns:
Augmented images and possibly flow, mask (if provided).
"""
# All photometric augmentation that could be applied individually per frame.
def potential_asymmetric_augmentations(images):
if probability_hue_shift > 0:
images = random_hue_shift(images, probability_hue_shift, max_delta_hue)
if probability_saturation > 0:
images = random_saturation(images, probability_saturation,
min_bound_saturation, max_bound_saturation)
if probability_brightness > 0:
images = random_brightness(images, probability_brightness,
max_delta_brightness)
if probability_contrast > 0:
images = random_contrast(images, probability_contrast, min_bound_contrast,
max_bound_contrast)
if probability_gaussian_noise > 0:
images = random_gaussian_noise(images, probability_gaussian_noise,
min_bound_gaussian_noise,
max_bound_gaussian_noise)
return images
perform_assymetric = tf.random.uniform([]) < probability_assymetric
def true_fn(images):
image_1, image_2 = tf.unstack(images)
image_1 = potential_asymmetric_augmentations(image_1)
image_2 = potential_asymmetric_augmentations(image_2)
return tf.stack([image_1, image_2])
def false_fn(images):
return images
images = tf.cond(perform_assymetric, lambda: true_fn(images),
lambda: false_fn(images))
# Photometric augmentations applied to all frames of a pair.
if probability_color_swap > 0:
images = random_color_swap(images, probability_color_swap)
# Photometric augmentations applied individually per image frame.
if probability_contrast_individual > 0:
images = random_contrast_individual(images, probability_contrast_individual,
min_bound_contrast_individual,
max_bound_contrast_individual)
if probability_brightness_individual > 0:
images = random_brightness_individual(images,
probability_brightness_individual,
max_delta_brightness_individual)
# Crop values to ensure values are within [0,1], some augmentations may
# violate this.
images = tf.clip_by_value(images, 0.0, 1.0)
# Apply special photometric augmentations.
if probability_eraser > 0:
images = random_eraser(
images,
min_size=min_size_eraser,
max_size=max_size_eraser,
probability=probability_eraser,
max_operations=max_operations_eraser,
probability_additional_operations=probability_eraser_additional_operations
)
return images
def project_distribution(supports, weights, target_support,
validate_args=False):
"""Projects a batch of (support, weights) onto target_support.
Based on equation (7) in (Bellemare et al., 2017):
https://arxiv.org/abs/1707.06887
In the rest of the comments we will refer to this equation simply as Eq7.
This code is not easy to digest, so we will use a running example to clarify
what is going on, with the following sample inputs:
* supports = [[0, 2, 4, 6, 8],
[1, 3, 4, 5, 6]]
* weights = [[0.1, 0.6, 0.1, 0.1, 0.1],
[0.1, 0.2, 0.5, 0.1, 0.1]]
* target_support = [4, 5, 6, 7, 8]
In the code below, comments preceded with 'Ex:' will be referencing the above
values.
Args:
supports: Tensor of shape (batch_size, num_dims) defining supports for the
distribution.
weights: Tensor of shape (batch_size, num_dims) defining weights on the
original support points. Although for the CategoricalDQN agent these
weights are probabilities, it is not required that they are.
target_support: Tensor of shape (num_dims) defining support of the projected
distribution. The values must be monotonically increasing. Vmin and Vmax
will be inferred from the first and last elements of this tensor,
respectively. The values in this tensor must be equally spaced.
validate_args: Whether we will verify the contents of the
target_support parameter.
Returns:
A Tensor of shape (batch_size, num_dims) with the projection of a batch of
(support, weights) onto target_support.
Raises:
ValueError: If target_support has no dimensions, or if shapes of supports,
weights, and target_support are incompatible.
"""
target_support_deltas = target_support[1:] - target_support[:-1]
# delta_z = `\Delta z` in Eq7.
delta_z = target_support_deltas[0]
validate_deps = []
supports.shape.assert_is_compatible_with(weights.shape)
supports[0].shape.assert_is_compatible_with(target_support.shape)
target_support.shape.assert_has_rank(1)
if validate_args:
# Assert that supports and weights have the same shapes.
validate_deps.append(
tf.Assert(
tf.reduce_all(tf.equal(tf.shape(supports), tf.shape(weights))),
[supports, weights]))
# Assert that elements of supports and target_support have the same shape.
validate_deps.append(
tf.Assert(
tf.reduce_all(
tf.equal(tf.shape(supports)[1], tf.shape(target_support))),
[supports, target_support]))
# Assert that target_support has a single dimension.
validate_deps.append(
tf.Assert(
tf.equal(tf.size(tf.shape(target_support)), 1), [target_support]))
# Assert that the target_support is monotonically increasing.
validate_deps.append(
tf.Assert(tf.reduce_all(target_support_deltas > 0), [target_support]))
# Assert that the values in target_support are equally spaced.
validate_deps.append(
tf.Assert(
tf.reduce_all(tf.equal(target_support_deltas, delta_z)),
[target_support]))
with tf.control_dependencies(validate_deps):
# Ex: `v_min, v_max = 4, 8`.
v_min, v_max = target_support[0], target_support[-1]
# Ex: `batch_size = 2`.
batch_size = tf.shape(supports)[0]
# `N` in Eq7.
# Ex: `num_dims = 5`.
num_dims = tf.shape(target_support)[0]
# clipped_support = `[\hat{T}_{z_j}]^{V_max}_{V_min}` in Eq7.
# Ex: `clipped_support = [[[ 4. 4. 4. 6. 8.]]
# [[ 4. 4. 4. 5. 6.]]]`.
clipped_support = tf.clip_by_value(supports, v_min, v_max)[:, None, :]
# Ex: `tiled_support = [[[[ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]
# [ 4. 4. 4. 6. 8.]]
# [[ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]
# [ 4. 4. 4. 5. 6.]]]]`.
tiled_support = tf.tile([clipped_support], [1, 1, num_dims, 1])
# Ex: `reshaped_target_support = [[[ 4.]
# [ 5.]
# [ 6.]
# [ 7.]
# [ 8.]]
# [[ 4.]
# [ 5.]
# [ 6.]
# [ 7.]
# [ 8.]]]`.
reshaped_target_support = tf.tile(target_support[:, None], [batch_size, 1])
reshaped_target_support = tf.reshape(reshaped_target_support,
[batch_size, num_dims, 1])
# numerator = `|clipped_support - z_i|` in Eq7.
# Ex: `numerator = [[[[ 0. 0. 0. 2. 4.]
# [ 1. 1. 1. 1. 3.]
# [ 2. 2. 2. 0. 2.]
# [ 3. 3. 3. 1. 1.]
# [ 4. 4. 4. 2. 0.]]
# [[ 0. 0. 0. 1. 2.]
# [ 1. 1. 1. 0. 1.]
# [ 2. 2. 2. 1. 0.]
# [ 3. 3. 3. 2. 1.]
# [ 4. 4. 4. 3. 2.]]]]`.
numerator = tf.abs(tiled_support - reshaped_target_support)
quotient = 1 - (numerator / delta_z)
# clipped_quotient = `[1 - numerator / (\Delta z)]_0^1` in Eq7.
# Ex: `clipped_quotient = [[[[ 1. 1. 1. 0. 0.]
# [ 0. 0. 0. 0. 0.]
# [ 0. 0. 0. 1. 0.]
# [ 0. 0. 0. 0. 0.]
# [ 0. 0. 0. 0. 1.]]
# [[ 1. 1. 1. 0. 0.]
# [ 0. 0. 0. 1. 0.]
# [ 0. 0. 0. 0. 1.]
# [ 0. 0. 0. 0. 0.]
# [ 0. 0. 0. 0. 0.]]]]`.
clipped_quotient = tf.clip_by_value(quotient, 0, 1)
# Ex: `weights = [[ 0.1 0.6 0.1 0.1 0.1]
# [ 0.1 0.2 0.5 0.1 0.1]]`.
weights = weights[:, None, :]
# inner_prod = `\sum_{j=0}^{N-1} clipped_quotient * p_j(x', \pi(x'))`
# in Eq7.
# Ex: `inner_prod = [[[[ 0.1 0.6 0.1 0. 0. ]
# [ 0. 0. 0. 0. 0. ]
# [ 0. 0. 0. 0.1 0. ]
# [ 0. 0. 0. 0. 0. ]
# [ 0. 0. 0. 0. 0.1]]
# [[ 0.1 0.2 0.5 0. 0. ]
# [ 0. 0. 0. 0.1 0. ]
# [ 0. 0. 0. 0. 0.1]
# [ 0. 0. 0. 0. 0. ]
# [ 0. 0. 0. 0. 0. ]]]]`.
inner_prod = clipped_quotient * weights
# Ex: `projection = [[ 0.8 0.0 0.1 0.0 0.1]
# [ 0.8 0.1 0.1 0.0 0.0]]`.
projection = tf.reduce_sum(inner_prod, 3)
projection = tf.reshape(projection, [batch_size, num_dims])
return projection
def wide_resnet_block(x,
depth,
stride,
weight_decay,
params=None,
moments=None,
use_project=False,
backprop_through_moments=True,
is_training=True,
use_bounded_activation=False):
"""Wide ResNet residual block."""
params_keys, params_vars = [], []
moments_keys, moments_vars = [], []
with tf.variable_scope('conv1'):
bn_1, bn_params, bn_moments = bn(
x,
params=params,
moments=moments,
is_training=is_training,
backprop_through_moments=backprop_through_moments)
params_keys.extend(bn_params.keys())
params_vars.extend(bn_params.values())
moments_keys.extend(bn_moments.keys())
moments_vars.extend(bn_moments.values())
out_1 = relu(bn_1, use_bounded_activation=use_bounded_activation)
h_1, conv_params = conv(
out_1, [3, 3], depth, stride, weight_decay, params=params)
params_keys.extend(conv_params.keys())
params_vars.extend(conv_params.values())
with tf.variable_scope('conv2'):
bn_2, bn_params, bn_moments = bn(
h_1,
params=params,
moments=moments,
is_training=is_training,
backprop_through_moments=backprop_through_moments)
params_keys.extend(bn_params.keys())
params_vars.extend(bn_params.values())
moments_keys.extend(bn_moments.keys())
moments_vars.extend(bn_moments.values())
out_2 = relu(bn_2, use_bounded_activation=use_bounded_activation)
h_2, conv_params = conv(
out_2, [3, 3],
depth,
stride=1,
weight_decay=weight_decay,
params=params)
params_keys.extend(conv_params.keys())
params_vars.extend(conv_params.values())
h = h_2
if use_bounded_activation:
h = tf.clip_by_value(h, -6, 6)
with tf.variable_scope('identity'):
if use_project:
with tf.variable_scope('projection_conv'):
x, conv_params = conv(
out_1, [1, 1], depth, stride, weight_decay, params=params)
params_keys.extend(conv_params.keys())
params_vars.extend(conv_params.values())
params = collections.OrderedDict(zip(params_keys, params_vars))
moments = collections.OrderedDict(zip(moments_keys, moments_vars))
if use_bounded_activation:
out = tf.clip_by_value(x + h, -6, 6)
else:
out = x + h
return out, params, moments
def bottleneck(x,
depth,
stride,
weight_decay,
params=None,
moments=None,
use_project=False,
backprop_through_moments=True,
is_training=True,
input_rate=1,
output_rate=1,
use_bounded_activation=False):
"""ResNet18 residual block."""
params_keys, params_vars = [], []
moments_keys, moments_vars = [], [] # means and vars of different layers.
with tf.variable_scope('conv1'):
h, conv_bn_params, conv_bn_moments = conv_bn(
x, [3, 3],
depth[0],
stride,
weight_decay,
params=params,
moments=moments,
is_training=is_training,
rate=input_rate,
backprop_through_moments=backprop_through_moments)
params_keys.extend(conv_bn_params.keys())
params_vars.extend(conv_bn_params.values())
moments_keys.extend(conv_bn_moments.keys())
moments_vars.extend(conv_bn_moments.values())
h = relu(h, use_bounded_activation=use_bounded_activation)
with tf.variable_scope('conv2'):
h, conv_bn_params, conv_bn_moments = conv_bn(
h, [3, 3],
depth[1],
stride=1,
weight_decay=weight_decay,
params=params,
moments=moments,
is_training=is_training,
rate=output_rate,
backprop_through_moments=backprop_through_moments)
if use_bounded_activation:
h = tf.clip_by_value(h, -6.0, 6.0)
params_keys.extend(conv_bn_params.keys())
params_vars.extend(conv_bn_params.values())
moments_keys.extend(conv_bn_moments.keys())
moments_vars.extend(conv_bn_moments.values())
with tf.variable_scope('identity'):
if use_project:
with tf.variable_scope('projection_conv'):
x, conv_bn_params, conv_bn_moments = conv_bn(
x, [1, 1],
depth[1],
stride,
weight_decay,
params=params,
moments=moments,
is_training=is_training,
rate=1,
backprop_through_moments=backprop_through_moments)
params_keys.extend(conv_bn_params.keys())
params_vars.extend(conv_bn_params.values())
moments_keys.extend(conv_bn_moments.keys())
moments_vars.extend(conv_bn_moments.values())
x = relu(x + h, use_bounded_activation=use_bounded_activation)
params = collections.OrderedDict(zip(params_keys, params_vars))
moments = collections.OrderedDict(zip(moments_keys, moments_vars))
return x, params, moments
def photometric_augmentation(images,
augment_color_swap=True,
augment_hue_shift=True,
augment_saturation=False,
augment_brightness=False,
augment_contrast=False,
augment_gaussian_noise=False,
augment_brightness_individual=False,
augment_contrast_individual=False,
max_delta_hue=0.5,
min_bound_saturation=0.8,
max_bound_saturation=1.2,
max_delta_brightness=0.1,
min_bound_contrast=0.8,
max_bound_contrast=1.2,
min_bound_gaussian_noise=0.0,
max_bound_gaussian_noise=0.02,
max_delta_brightness_individual=0.02,
min_bound_contrast_individual=0.95,
max_bound_contrast_individual=1.05):
"""Applies photometric augmentations to an image pair."""
# Randomly permute colors by rolling and reversing.
# This covers all permutations.
if augment_color_swap:
r = tf.random.uniform([], maxval=3, dtype=tf.int32)
images = tf.roll(images, r, axis=-1)
r = tf.equal(tf.random.uniform([], maxval=2, dtype=tf.int32), 1)
images = tf.cond(pred=r,
true_fn=lambda: tf.reverse(images, axis=[-1]),
false_fn=lambda: images)
if augment_hue_shift:
images = tf.image.random_hue(images, max_delta_hue)
if augment_saturation:
images = tf.image.random_saturation(
images, min_bound_saturation, max_bound_saturation)
if augment_brightness:
images = tf.image.random_brightness(images, max_delta_brightness)
if augment_contrast:
images = tf.image.random_contrast(
images, min_bound_contrast, max_bound_contrast)
if augment_gaussian_noise:
sigma = tf.random.uniform([],
minval=min_bound_gaussian_noise,
maxval=max_bound_gaussian_noise,
dtype=tf.float32)
noise = tf.random.normal(
tf.shape(input=images), stddev=sigma, dtype=tf.float32)
images = images + noise
# perform relative photometric augmentation (individually per image)
image_1, image_2 = tf.unstack(images)
if augment_brightness_individual:
image_1 = tf.image.random_contrast(
image_1, min_bound_contrast_individual, max_bound_contrast_individual)
image_2 = tf.image.random_contrast(
image_2, min_bound_contrast_individual, max_bound_contrast_individual)
if augment_contrast_individual:
image_1 = tf.image.random_brightness(
image_1, max_delta_brightness_individual)
image_2 = tf.image.random_brightness(
image_2, max_delta_brightness_individual)
# crop values to ensure values in [0,1] (some augmentations can violate this)
image_1 = tf.clip_by_value(image_1, 0.0, 1.0)
image_2 = tf.clip_by_value(image_2, 0.0, 1.0)
return tf.stack([image_1, image_2])
