def compute_gradients(self,
loss,
var_list=None,
gate_gradients=GATE_OP,
aggregation_method=None,
colocate_gradients_with_ops=False,
grad_loss=None):
if self._use_collection_losses:
loss = tf.get_collection(PCGRAD_LOSSES_COLLECTION)
assert isinstance(loss,
list), "The loss is not a list: %s" % type(loss)
num_tasks = len(loss)
loss = tf.stack(loss)
tf.random.shuffle(loss)
# Compute per-task gradients.
def compute_per_task_grads(task):
grad_list = []
for grad in tf.gradients(task, var_list):
if grad is None:
continue
grad_list.append(tf.reshape(grad, [
-1,
]))
return tf.concat(grad_list, axis=0)
grads_task = tf.vectorized_map(compute_per_task_grads, loss)
# Compute gradient projections.
def proj_grad(grad_task):
for k in range(num_tasks):
inner_product = tf.reduce_sum(grad_task * grads_task[k])
proj_direction = inner_product / tf.reduce_sum(
grads_task[k] * grads_task[k])
grad_task = grad_task - tf.minimum(proj_direction,
0.) * grads_task[k]
return grad_task
proj_grads_flatten = tf.vectorized_map(proj_grad, grads_task)
# Unpack flattened projected gradients back to their original shapes.
proj_grads = []
for j in range(num_tasks):
start_idx = 0
for idx, var in enumerate(var_list):
grad_shape = var.get_shape()
flatten_dim = np.prod([
grad_shape.dims[i].value
for i in range(len(grad_shape.dims))
])
proj_grad = proj_grads_flatten[j][start_idx:start_idx +
flatten_dim]
proj_grad = tf.reshape(proj_grad, grad_shape)
if len(proj_grads) < len(var_list):
proj_grads.append(proj_grad)
else:
proj_grads[idx] += proj_grad
start_idx += flatten_dim
grads_and_vars = list(zip(proj_grads, var_list))
return grads_and_vars
def _compute_projected_grads(self, pcgrad_vars, loss):
if not pcgrad_vars:
return []
num_tasks = len(loss)
loss = tf.stack(loss)
# Compute per-task gradients.
def compute_per_task_grads(task):
grad_list = []
original_pcgrad_grads_vars = self._optimizer.compute_gradients(
task, pcgrad_vars)
original_grads = [
grad_var[0] for grad_var in original_pcgrad_grads_vars
]
for grad in original_grads:
if grad is None:
continue
grad_list.append(tf.reshape(grad, [
-1,
]))
return tf.concat(grad_list, axis=0)
grads_task = tf.vectorized_map(compute_per_task_grads, loss)
# Compute gradient projections.
def proj_grad(grad_task):
for k in range(num_tasks):
inner_product = tf.reduce_sum(grad_task * grads_task[k])
proj_direction = inner_product / tf.reduce_sum(
grads_task[k] * grads_task[k])
grad_task = grad_task - tf.minimum(proj_direction,
0.) * grads_task[k]
return grad_task
proj_grads_flatten = tf.vectorized_map(proj_grad, grads_task)
# Unpack flattened projected gradients back to their original shapes.
proj_grads = []
for j in range(num_tasks):
start_idx = 0
for idx, var in enumerate(pcgrad_vars):
grad_shape = var.get_shape()
flatten_dim = np.prod([
grad_shape.dims[i].value
for i in range(len(grad_shape.dims))
])
proj_grad = proj_grads_flatten[j][start_idx:start_idx +
flatten_dim]
proj_grad = tf.reshape(proj_grad, grad_shape)
if len(proj_grads) < len(pcgrad_vars):
proj_grads.append(proj_grad)
else:
proj_grads[idx] += proj_grad
start_idx += flatten_dim
return proj_grads
def pfor_map_fn(fn, elems):
"""Provides a map_fn-like wrapper around pfor.
TODO(T2R_CONTRIBUTORS): For now, fn's conv filters should be loop invariant.
TODO(T2R_CONTRIBUTORS): Using `AssignSubvariableOp` in fn will cause errors.
Args:
fn: The callable to be performed. It accepts one argument, which will have
the same (possibly nested) structure as `elems`. It should output a
tensor or structure of tensors.
elems: A tensor or (possibly nested) structure of tensors, each of which
will be unpacked along their first dimension. The nested sequence of
resulting slices will be applied to `fn`.
Returns:
A tensor or (possibly nested) structure of tensors, resulting from stacking
the outputs of `fn` on the unpacked `elems`.
"""
def loop_fn(i):
"""Calls fn on the i'th entry of each tensor in `elems`."""
gathered_elems = tf.nest.map_structure(lambda x: tf.gather(x, i),
elems)
return fn(gathered_elems)
batch_size = tf.shape(tf.nest.flatten(elems)[0])[0]
return tf.vectorized_map(loop_fn, batch_size)
def batch_image_preprocess(raw_images,
image_size: Union[int, Tuple[int, int]],
mean_rgb,
stddev_rgb,
batch_size: int = None):
"""Preprocess batched images for inference.
Args:
raw_images: a list of images, each image can be a tensor or a numpy arary.
image_size: single integer of image size for square image or tuple of two
integers, in the format of (image_height, image_width).
mean_rgb: Mean value of RGB, can be a list of float or a float value.
stddev_rgb: Standard deviation of RGB, can be a list of float or a float
value.
batch_size: if None, use map_fn to deal with dynamic batch size.
Returns:
(image, scale): a tuple of processed images and scales.
"""
if not batch_size:
# map_fn is a little bit slower due to some extra overhead.
# map_fn -> vectorized_map (fully parallelizes the batch).
map_fn = functools.partial(
image_preprocess,
image_size=image_size,
mean_rgb=mean_rgb,
stddev_rgb=stddev_rgb)
images, scales = tf.vectorized_map(map_fn, raw_images)
images = tf.stop_gradient(tf.cast(images, tf.float32))
scales = tf.stop_gradient(tf.cast(scales, tf.float32))
return (images, scales)
# If batch size is known, use a simple loop.
scales, images = [], []
for i in range(batch_size):
image, scale = image_preprocess(raw_images[i], image_size, mean_rgb,
stddev_rgb)
scales.append(scale)
images.append(image)
images = tf.stack(images)
scales = tf.stack(scales)
return (images, scales)
def resnet_v2_152(inputs,
y,
num_classes=None,
is_training=True,
global_pool=True,
output_stride=None,
spatial_squeeze=True,
reuse=None,
scope='resnet_v2_152'):
"""ResNet-152 model of [1]. See resnet_v2() for arg and return description."""
blocks = [
resnet_v2_block('block1', base_depth=64, num_units=3, stride=2),
resnet_v2_block('block2', base_depth=128, num_units=8, stride=2),
resnet_v2_block('block3', base_depth=256, num_units=36, stride=2),
resnet_v2_block('block4', base_depth=512, num_units=3, stride=1),
]
net,outputs,scopes =resnet_v2(inputs, blocks, num_classes, is_training=is_training,
global_pool=global_pool, output_stride=output_stride,
include_root_block=True, spatial_squeeze=spatial_squeeze,
reuse=reuse, scope=scope)
net = tf.squeeze(net, [1, 2], name="squzzezd")
def fn(args):
y, index = args
return tf.gather(y, index)
_, indexs = tf.math.top_k(net, 5)
acc_array = tf.vectorized_map(fn, (y, indexs))
top_accuracy = tf.reduce_sum(acc_array, name="top_accuracy")
with tf.variable_scope(scopes[-1]):
loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=net)
loss = tf.reduce_mean(loss)
outputs[-1]=loss
return loss,outputs,scopes
def compute_gradients(self,
loss,
var_list,
gate_gradients=GATE_OP,
aggregation_method=None,
colocate_gradients_with_ops=False,
grad_loss=None,
gradient_tape=None):
if callable(loss):
# TF is running in Eager mode
raise NotImplementedError(
'Vectorized optimizer unavailable for TF2.')
else:
# TF is running in graph mode, check we did not receive a gradient tape.
if gradient_tape:
raise ValueError(
'When in graph mode, a tape should not be passed.')
batch_size = tf.shape(input=loss)[0]
if self._num_microbatches is None:
self._num_microbatches = batch_size
# Note: it would be closer to the correct i.i.d. sampling of records if
# we sampled each microbatch from the appropriate binomial distribution,
# although that still wouldn't be quite correct because it would be
# sampling from the dataset without replacement.
microbatch_losses = tf.reshape(loss,
[self._num_microbatches, -1])
if var_list is None:
var_list = (tf.trainable_variables() + tf.get_collection(
tf.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
def process_microbatch(microbatch_loss):
"""Compute clipped grads for one microbatch."""
microbatch_loss = tf.reduce_mean(
input_tensor=microbatch_loss)
grads, _ = zip(
*super(DPOptimizerClass, self).compute_gradients(
microbatch_loss, var_list, gate_gradients,
aggregation_method, colocate_gradients_with_ops,
grad_loss))
grads_list = [
g if g is not None else tf.zeros_like(v)
for (g, v) in zip(list(grads), var_list)
]
# Clip gradients to have L2 norm of l2_norm_clip.
# Here, we use TF primitives rather than the built-in
# tf.clip_by_global_norm() so that operations can be vectorized
# across microbatches.
grads_flat = tf.nest.flatten(grads_list)
squared_l2_norms = [
tf.reduce_sum(input_tensor=tf.square(g))
for g in grads_flat
]
global_norm = tf.sqrt(tf.add_n(squared_l2_norms))
div = tf.maximum(global_norm / self._l2_norm_clip, 1.)
clipped_flat = [g / div for g in grads_flat]
clipped_grads = tf.nest.pack_sequence_as(
grads_list, clipped_flat)
return clipped_grads
clipped_grads = tf.vectorized_map(process_microbatch,
microbatch_losses)
def reduce_noise_normalize_batch(stacked_grads):
summed_grads = tf.reduce_sum(input_tensor=stacked_grads,
axis=0)
noise_stddev = self._l2_norm_clip * self._noise_multiplier
noise = tf.random.normal(tf.shape(input=summed_grads),
stddev=noise_stddev)
noised_grads = summed_grads + noise
return noised_grads / tf.cast(self._num_microbatches,
tf.float32)
final_grads = tf.nest.map_structure(
reduce_noise_normalize_batch, clipped_grads)
return list(zip(final_grads, var_list))
def vgg_19(inputs,
y,
num_classes=1000,
is_training=True,
dropout_keep_prob=0.5,
spatial_squeeze=True,
reuse=None,
scope='vgg_19',
fc_conv_padding='VALID',
global_pool=False):
"""Oxford Net VGG 19-Layers version E Example.
Note: All the fully_connected layers have been transformed to conv2d layers.
To use in classification mode, resize input to 224x224.
Args:
inputs: a tensor of size [batch_size, height, width, channels].
num_classes: number of predicted classes. If 0 or None, the logits layer is
omitted and the input features to the logits layer are returned instead.
is_training: whether or not the model is being trained.
dropout_keep_prob: the probability that activations are kept in the dropout
layers during training.
spatial_squeeze: whether or not should squeeze the spatial dimensions of the
outputs. Useful to remove unnecessary dimensions for classification.
reuse: whether or not the network and its variables should be reused. To be
able to reuse 'scope' must be given.
scope: Optional scope for the variables.
fc_conv_padding: the type of padding to use for the fully connected layer
that is implemented as a convolutional layer. Use 'SAME' padding if you
are applying the network in a fully convolutional manner and want to
get a prediction map downsampled by a factor of 32 as an output.
Otherwise, the output prediction map will be (input / 32) - 6 in case of
'VALID' padding.
global_pool: Optional boolean flag. If True, the input to the classification
layer is avgpooled to size 1x1, for any input size. (This is not part
of the original VGG architecture.)
Returns:
net: the output of the logits layer (if num_classes is a non-zero integer),
or the non-dropped-out input to the logits layer (if num_classes is 0 or
None).
end_points: a dict of tensors with intermediate activations.
"""
scopes = []
outputs= []
if True:
tf.get_variable_scope()._reuse=tf.AUTO_REUSE
scope_name = tf.get_variable_scope().name
end_points_collection = tf.get_variable_scope().name + '_end_points'
# Collect outputs for conv2d, fully_connected and max_pool2d.
with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
outputs_collections=end_points_collection):
net = slim.repeat(inputs, 2, slim.conv2d, 64, [3, 3], scope='conv1')
scopes.append('conv1')
outputs.append(net)
net = slim.max_pool2d(net, [2, 2], scope='pool1')
scopes.append('pool1')
outputs.append(net)
net = slim.repeat(net, 2, slim.conv2d, 128, [3, 3], scope='conv2')
scopes.append('conv2')
outputs.append(net)
net = slim.max_pool2d(net, [2, 2], scope='pool2')
scopes.append('pool2')
outputs.append(net)
net = slim.repeat(net, 4, slim.conv2d, 256, [3, 3], scope='conv3')
scopes.append('conv3')
outputs.append(net)
net = slim.max_pool2d(net, [2, 2], scope='pool3')
scopes.append('pool3')
outputs.append(net)
net = slim.repeat(net, 4, slim.conv2d, 512, [3, 3], scope='conv4')
scopes.append('conv4')
outputs.append(net)
net = slim.max_pool2d(net, [2, 2], scope='pool4')
scopes.append('pool4')
outputs.append(net)
net = slim.repeat(net, 4, slim.conv2d, 512, [3, 3], scope='conv5')
scopes.append('conv5')
outputs.append(net)
net = slim.max_pool2d(net, [2, 2], scope='pool5')
scopes.append('pool5')
outputs.append(net)
# Use conv2d instead of fully_connected layers.
net = slim.conv2d(net, 4096, [7, 7], padding=fc_conv_padding, scope='fc6')
scopes.append('fc6')
outputs.append(net)
net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
scope='dropout6')
scopes.append('dropout6')
outputs.append(net)
net = slim.conv2d(net, 4096, [1, 1], scope='fc7')
scopes.append('fc7')
outputs.append(net)
net = slim.conv2d(net, 4096, [1, 1], scope='fc8')
scopes.append('fc8')
outputs.append(net)
net = slim.conv2d(net, 4096, [1, 1], scope='fc9')
scopes.append('fc9')
outputs.append(net)
net = slim.conv2d(net, 4096, [1, 1], scope='fc10')
scopes.append('fc10')
outputs.append(net)
# Convert end_points_collection into a end_point dict.
if num_classes:
net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
scope='dropout10')
scopes.append('dropout10')
outputs.append(net)
net = slim.conv2d(net, num_classes, [1, 1],
activation_fn=None,
normalizer_fn=None,
scope='fc11')
with tf.variable_scope("fc11"):
net = tf.squeeze(net, [1, 2], name="squzzezd")
_, indexs = tf.math.top_k(net,5)
def fn(args):
y,index = args
return tf.gather(y,index)
acc_array = tf.vectorized_map(fn,(y,indexs))
top_accuracy = tf.reduce_sum(acc_array,name="top_accuracy")
loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=net)
loss = tf.reduce_mean(loss)
scopes.append('fc11')
outputs.append(loss)
return loss, outputs,scopes