def prep_fn(image, label):
"""Image preprocessing function."""
if config.roll_pixels:
image = tf.roll(image, config.roll_pixels, -2)
if is_training:
image = tf.image.random_flip_left_right(image)
image = tf.pad(image, [[4, 4], [4, 4], [0, 0]])
image = tf.image.random_crop(image, CIFAR_SHAPE)
image = tf.image.convert_image_dtype(image, tf.float32)
return image, label
def _forward(self, x):
ndims = prefer_static.rank(x)
indices = prefer_static.reshape(prefer_static.add(self.axis, ndims),
shape=[-1, 1])
return tf.pad(
x,
paddings=prefer_static.tensor_scatter_nd_update(
prefer_static.zeros([ndims, 2], dtype=tf.int32),
indices, self.paddings),
mode=self.mode,
constant_values=prefer_static.cast(self.constant_values, dtype=x.dtype))
def loop_body(args):
prob, length = args
prob = prob[:length]
overflow = tf.math.maximum(
1 - tf.reduce_sum(prob, keepdims=True), 0.)
prob = tf.concat([prob, overflow], axis=0)
cdf = range_coding_ops.pmf_to_quantized_cdf(
prob, precision=self.range_coder_precision)
return tf.pad(cdf, [[0, max_length - length]],
mode="CONSTANT",
constant_values=0)
def _instruction(self, instruction, length):
"""Processing of the language instructions."""
embedding = self._text_embedder(instruction)
# Pad to make sure there is at least one output.
padding = tf.cast(tf.equal(tf.shape(embedding)[1], 0), dtype=tf.int32)
embedding = tf.pad(embedding, [[0, 0], [0, padding], [0, 0]])
output = self._instruction_rnn_layer(embedding)
output = tf.reverse_sequence(output, length, seq_axis=1)[:, 0]
output = tf.keras.layers.Flatten()(output)
return output
def pad_images_if_uneven(images):
# If num partitions is larger than 1, then pad the input so that
# it passes JAX's ragged partition checking.
# Otherwise it throws an error when trying to partition second
# dimension (150) into 4 parts.
# Image shape [local_device, B, H, W, C] --> # [B, H, W, C]
padding_size = get_spmd_image_padding_size(params, images.shape[2:])
if padding_size:
padding_spec = tf.constant([[0, 0], [0, 0], [0, padding_size],
[0, 0], [0, 0]])
images = tf.pad(images, padding_spec)
return images
def tile_for_batch(x, batch_shape):
x_shape = tf.shape(x)
x_ndims = tf.rank(x)
x = tf.reshape(x, shape=tf.concat([
x_shape[:-1],
tf.ones_like(batch_shape),
x_shape[-1:],
], axis=0))
return tf.tile(x, multiples=tf.pad(
batch_shape,
paddings=[[x_ndims - 1, 1]],
constant_values=1))
def preprocess_data(images, labels, is_training):
"""CIFAR data preprocessing"""
images = tf.image.convert_image_dtype(images, tf.float32)
if is_training:
crop_padding = 4
images = tf.pad(images, [[crop_padding, crop_padding],
[crop_padding, crop_padding], [0, 0]],
'REFLECT')
images = tf.image.random_crop(images, [32, 32, 3])
images = tf.image.random_flip_left_right(images)
return images, labels
def pad_shape_with_ones(x, ndims, start=-1):
"""Maybe add `ndims` ones to `x.shape` starting at `start`.
If `ndims` is zero, this is a no-op; otherwise, we will create and return a
new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
right side. If the shape of `x` is known statically, the shape of the return
value will be as well.
Args:
x: The `Tensor` we'll return a reshaping of.
ndims: Python `integer` number of ones to pad onto `x.shape`.
start: Python `integer` specifying where to start padding with ones. Must
be a negative integer. For instance, a value of `-1` means to pad at the
end of the shape. Default value: `-1`.
Returns:
If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
with `ndims` ones concatenated on the right side. If possible, returns a
`Tensor` whose shape is known statically.
Raises:
ValueError: if `ndims` is not a Python `integer` greater than or equal to
zero.
"""
if not (isinstance(ndims, int) and ndims >= 0):
raise ValueError(
'`ndims` must be a Python `integer` greater than zero. Got: {}'
.format(ndims))
if not (isinstance(start, int) and start <= -1):
raise ValueError(
'`start` must be a Python `integer` less than zero. Got: {}'
.format(start))
if ndims == 0:
return x
x = tf.convert_to_tensor(value=x)
original_shape = x.shape
rank = tf.rank(x)
first_shape = tf.shape(x)[:rank + start + 1]
second_shape = tf.shape(x)[rank + start + 1:]
new_shape = tf.pad(first_shape, paddings=[[0, ndims]], constant_values=1)
new_shape = tf.concat([new_shape, second_shape], axis=0)
x = tf.reshape(x, new_shape)
if start == -1:
tensorshape_util.set_shape(
x, tensorshape_util.concatenate(original_shape, [1] * ndims))
elif tensorshape_util.rank(original_shape) is not None:
original_ndims = tensorshape_util.rank(original_shape)
new_shape = tensorshape_util.concatenate(
original_shape[:original_ndims + start + 1],
tensorshape_util.concatenate(
[1] * ndims,
original_shape[original_ndims + start + 1:]))
tensorshape_util.set_shape(x, new_shape)
return x
def preprocess_data(images, labels, is_training):
"""CIFAR data preprocessing"""
images = tf.image.convert_image_dtype(images, tf.float32)
if is_training:
crop_padding = 4
images = tf.pad(images, [[crop_padding, crop_padding],
[crop_padding, crop_padding], [0, 0]], 'REFLECT')
images = tf.image.random_crop(images, [32, 32, 3])
images = tf.image.random_flip_left_right(images)
if FLAGS.distort_color:
images = color_distortion(images, s=1.0)
return {'input_1': images, 'label': labels}
def testPadWithOneReflectionIsCorrect(self):
"""Tests that pad_reflecting(p) matches tf.pad(p) when p is small."""
for _ in range(4):
n = int(np.ceil(np.random.uniform() * 8)) + 1
x = np.random.uniform(size=(n, n, n))
padding_below = int(np.round(np.random.uniform() * (n - 1)))
padding_above = int(np.round(np.random.uniform() * (n - 1)))
axis = int(np.floor(np.random.uniform() * 3.))
if axis == 0:
reference = tf.pad(
x, [[padding_below, padding_above], [0, 0], [0, 0]], mode='REFLECT')
elif axis == 1:
reference = tf.pad(
x, [[0, 0], [padding_below, padding_above], [0, 0]], mode='REFLECT')
elif axis == 2:
reference = tf.pad(
x, [[0, 0], [0, 0], [padding_below, padding_above]], mode='REFLECT')
result = wavelet.pad_reflecting(x, padding_below, padding_above, axis)
self.assertAllEqual(result.shape, reference.shape)
self.assertAllEqual(result, reference)
def preprocess(image, label):
"""Image preprocessing function."""
if data_augmentation and split == tfds.Split.TRAIN:
image = tf.image.random_flip_left_right(image)
image = tf.pad(image, [[4, 4], [4, 4], [0, 0]])
image = tf.image.random_crop(image, image_shape)
image = tf.image.convert_image_dtype(image, tf.float32)
mean = tf.constant([0.4914, 0.4822, 0.4465])
std = tf.constant([0.2023, 0.1994, 0.2010])
image = (image - mean) / std
label = tf.cast(label, tf.float32)
return image, label
def _generic_neighbors_op(self, segment_fn, node_data, edge_mask=None):
"Return the maxima of node_data of adjacent nodes (not including self)."
with tf.name_scope(self.scope):
assert node_data.shape == (self.n, self.max_order)
node_data_f = tf.reshape(node_data, (self.n * self.max_order, ))
k = self.v_tot_size * 2
edge_data = tf.gather(node_data_f, self.v_edge_starts_global[:k])
if edge_mask is not None:
edge_data = tf.cast(edge_mask, node_data.dtype) * edge_data
node_out = segment_fn(edge_data, self.v_edge_ends_global[:k])
pad_out = self.n * self.max_order - tf.shape(node_out)[0]
return tf.reshape(tf.pad(node_out, [[0, pad_out]]),
(self.n, self.max_order))
def cutout(image: tf.Tensor, pad_size: int, replace: int = 0) -> tf.Tensor:
"""Apply cutout (https://arxiv.org/abs/1708.04552) to image.
This operation applies a (2*pad_size x 2*pad_size) mask of zeros to
a random location within `img`. The pixel values filled in will be of the
value `replace`. The located where the mask will be applied is randomly
chosen uniformly over the whole image.
Args:
image: An image Tensor of type uint8.
pad_size: Specifies how big the zero mask that will be generated is that
is applied to the image. The mask will be of size
(2*pad_size x 2*pad_size).
replace: What pixel value to fill in the image in the area that has
the cutout mask applied to it.
Returns:
An image Tensor that is of type uint8.
"""
image_height = tf.shape(image)[0]
image_width = tf.shape(image)[1]
# Sample the center location in the image where the zero mask will be applied.
cutout_center_height = tf.random.uniform(shape=[],
minval=0,
maxval=image_height,
dtype=tf.int32)
cutout_center_width = tf.random.uniform(shape=[],
minval=0,
maxval=image_width,
dtype=tf.int32)
lower_pad = tf.maximum(0, cutout_center_height - pad_size)
upper_pad = tf.maximum(0, image_height - cutout_center_height - pad_size)
left_pad = tf.maximum(0, cutout_center_width - pad_size)
right_pad = tf.maximum(0, image_width - cutout_center_width - pad_size)
cutout_shape = [
image_height - (lower_pad + upper_pad),
image_width - (left_pad + right_pad)
]
padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]
mask = tf.pad(tf.zeros(cutout_shape, dtype=image.dtype),
padding_dims,
constant_values=1)
mask = tf.expand_dims(mask, -1)
mask = tf.tile(mask, [1, 1, 3])
image = tf.where(tf.equal(mask, 0),
tf.ones_like(image, dtype=image.dtype) * replace, image)
return image
def _apply_mixed_term_explicitly(values_with_boundaries, mixed_term, delta_t,
dim1, dim2, has_default_lower_boundary,
has_default_upper_boundary, n_dims):
"""Applies mixed term explicitly."""
(mixed_term_pp, mixed_term_pm, mixed_term_mp, mixed_term_mm) = mixed_term
batch_rank = values_with_boundaries.shape.rank - n_dims
# Below we multiply the mixed terms by inner value grid "shifted" diagonally.
# With Robin boundary conditions, this shift is done by restoring the
# boundaries (values_with_boundaries already have them restored) and then
# slicing. E.g. in 2D, values_inner = values_with_boundaries[1:-1, 1:-1],
# and an example of a diagonally-shifted slice is
# values_with_boundaries[:-2, 2:], which gets multiplied by mixed_term_mp.
# However, with default boundaries, there's no boundary to restore, and the
# diagonal shift go out of bounds. Since the mixed terms on the default
# boundaries are zero, the "out-of-bounds" values are irrelevant. However, to
# avoid going out of bounds and get the shapes right, we need to pad these
# boundaries.
paddings = batch_rank * [[0, 0]]
for dim in range(n_dims):
lower = 1 if has_default_lower_boundary[dim] else 0
upper = 1 if has_default_upper_boundary[dim] else 0
paddings += [[lower, upper]]
# Pad default boundaries with zeros
values_with_boundaries = tf.pad(values_with_boundaries, paddings=paddings)
def create_trimming_shifts(dim1_shift, dim2_shift):
# See _trim_boundaries. We need to apply shifts to dimensions dim1 and dim2.
shifts = [0] * n_dims
shifts[dim1] = dim1_shift
shifts[dim2] = dim2_shift
return shifts
values_pp = _trim_boundaries(values_with_boundaries,
from_dim=batch_rank,
shifts=create_trimming_shifts(1, 1))
values_mm = _trim_boundaries(values_with_boundaries,
from_dim=batch_rank,
shifts=create_trimming_shifts(-1, -1))
values_pm = _trim_boundaries(values_with_boundaries,
from_dim=batch_rank,
shifts=create_trimming_shifts(1, -1))
values_mp = _trim_boundaries(values_with_boundaries,
from_dim=batch_rank,
shifts=create_trimming_shifts(-1, 1))
return (mixed_term_mm * values_mm + mixed_term_mp * values_mp +
mixed_term_pm * values_pm + mixed_term_pp * values_pp) * delta_t
def _pad(batch):
"""Helper function to pad nested data within each batch elements."""
padded_dict_batch = {}
if isinstance(batch, dict):
for key, value in batch.items():
padded_dict_batch[key] = _pad(value)
return padded_dict_batch
rank = len(batch.shape)
assert rank > 0
missing_count = (self.padded_batch_size -
self.get_real_batch_size(batch))
padding = backend.stack([[0, missing_count]] + [[0, 0]] * (rank - 1))
return tf.pad(batch, padding, 'constant')
def _per_timestep_call(self, current_outputs, new_inputs, length, timestep,
**kwargs):
"""Returns Tensor of shape [..., timestep+1, vocab_size].
Args:
current_outputs: Tensor of shape [..., timestep, vocab_size], the so-far
generated sequence Tensor.
new_inputs: Tensor of shape [..., vocab_size], the new input to generate
its output given current_outputs.
length: Length of final desired sequence.
timestep: Current timestep.
**kwargs: Optional keyword arguments to layer.
"""
inputs = tf.concat([current_outputs, new_inputs[..., tf.newaxis, :]],
axis=-2)
# TODO(trandustin): To handle variable lengths, extend MADE to subset its
# input and output layer weights rather than pad inputs.
batch_ndims = inputs.shape.ndims - 2
padded_inputs = tf.pad(inputs,
paddings=[[0, 0]] * batch_ndims +
[[0, length - timestep - 1], [0, 0]])
net = self.layer(padded_inputs, **kwargs)
if net.shape[-1] == 2 * self.vocab_size:
loc, scale = tf.split(net, 2, axis=-1)
loc = loc[..., :(timestep + 1), :]
loc = tf.cast(utils.one_hot_argmax(loc, self.temperature),
inputs.dtype)
scale = scale[..., :(timestep + 1), :]
scale = tf.cast(utils.one_hot_argmax(scale, self.temperature),
inputs.dtype)
inverse_scale = utils.multiplicative_inverse(
scale, self.vocab_size)
shifted_inputs = utils.one_hot_minus(inputs, loc)
new_outputs = utils.one_hot_multiply(shifted_inputs, inverse_scale)
elif net.shape[-1] == self.vocab_size:
loc = net
loc = loc[..., :(timestep + 1), :]
loc = tf.cast(utils.one_hot_argmax(loc, self.temperature),
inputs.dtype)
new_outputs = utils.one_hot_minus(inputs, loc)
else:
raise ValueError(
'Output of layer does not have compatible dimensions.')
outputs = tf.concat([current_outputs, new_outputs[..., -1:, :]],
axis=-2)
if not tf.executing_eagerly():
outputs.set_shape([None] * batch_ndims +
[timestep + 1, self.vocab_size])
return outputs
def __init__(self, params):
"""Constructor for BertModel.
Args:
params: `BigBirdConfig` dictionary.
"""
self.params = copy.deepcopy(params)
self.scope = params["scope"]
super(BertModel, self).__init__(name=self.scope)
# validate params
self.pad = lambda x: x
if params["max_encoder_length"] <= 512:
logging.info("Switching to full attention for short sequences")
self.params["attention_type"] = "original_full"
if self.params["attention_type"] == "simulated_sparse" or self.params[
"attention_type"] == "block_sparse":
if params["max_encoder_length"] % params["block_size"]:
logging.info(
"Expand max_encoder_length to next multiple of block_size")
self.params["max_encoder_length"] = (
params["max_encoder_length"] // params["block_size"] +
1) * params["block_size"]
pad_size = self.params["max_encoder_length"] - params[
"max_encoder_length"]
paddings = [[0, 0], [0, pad_size]]
self.pad = lambda x: tf.pad(x, paddings)
with tf.compat.v1.variable_scope(self.scope,
reuse=tf.compat.v1.AUTO_REUSE):
self.embeder = utils.EmbeddingLayer(
vocab_size=self.params["vocab_size"],
emb_dim=self.params["hidden_size"],
initializer=utils.create_initializer(
self.params["initializer_range"]),
scale_emb=self.params["rescale_embedding"],
use_token_type=True,
num_token_types=self.params["type_vocab_size"],
use_position_embeddings=True,
max_position_embeddings=self.params["max_position_embeddings"],
dropout_prob=self.params["hidden_dropout_prob"])
self.encoder = encoder.EncoderStack(self.params)
self.pooler = utils.SimpleDenseLayer(
input_size=self.params["hidden_size"],
output_size=self.params["hidden_size"],
initializer=utils.create_initializer(
self.params["initializer_range"]),
activation=tf.tanh,
name="pooler/dense")
def preprocess_data_with_id(data, is_training):
"""CIFAR data preprocessing when image ids are included in the data loader"""
image = data['image']
image = tf.image.convert_image_dtype(image, tf.float32)
if is_training:
crop_padding = 4
image = tf.pad(image, [[crop_padding, crop_padding],
[crop_padding, crop_padding], [0, 0]],
'REFLECT')
image = tf.image.random_crop(image, [32, 32, 3])
image = tf.image.random_flip_left_right(image)
else:
image = tf.image.resize_with_crop_or_pad(image, 32, 32) # central crop
return data['id'], image, data['label']
def frame_audio(self, audio, hop_length=1024, center=True):
"""Slice audio into frames for crepe."""
# Pad so that frames are centered around their timestamps.
# (i.e. first frame is zero centered).
pad = int(self.frame_length / 2)
audio = tf.pad(audio, ((0, 0), (pad, pad))) if center else audio
frames = tf.signal.frame(audio,
frame_length=self.frame_length,
frame_step=hop_length)
# Normalize each frame -- this is expected by the model.
mean, var = tf.nn.moments(frames, [-1], keepdims=True)
frames -= mean
frames /= (var**0.5 + 1e-5)
return frames
def call(self, x):
shape = x.shape
rank = len(shape)
dim = self.dimension + 1
# Assume 1 batch_dim.
index = [0] * len(self.resolution)
y = x
paddings = np.zeros((rank, 2), dtype=np.int32)
paddings[dim, 0] = 1
y = tf.pad(y, paddings)
rem_dims = rank - 1 - len(index[:dim])
slice_inds = [0] + index[:dim] + [0] * rem_dims
return tf.slice(y, slice_inds, shape)
def loss_gradient_fn(objective, constraints, proxy_constraints, state):
"""Evaluates the loss for the current internal state."""
# Make sure that the objective and proxy constraints have the same dtype.
if (constraints.dtype.base_dtype != objective.dtype.base_dtype
or proxy_constraints.dtype.base_dtype !=
objective.dtype.base_dtype):
raise TypeError(
"objective, constraints and proxy_constraints must all "
"have the same dtypes")
distribution = self._distribution(state)
zero_and_constraints = tf.pad(constraints, [[1, 0]])
# The proxy-Lagrangian is defined as:
# objective_and_proxy_constraints = tf.concat(
# (tf.expand_dims(objective, 0), proxy_constraints), axis=0)
# output = tf.tensordot(
# distribution, objective_and_proxy_constraints, axes=1)
# when updating the model parameters, and as:
# output = tf.tensordot(
# distribution, zero_and_constraints, axes=1)
# when updating the internal state.
#
# However, while these quantities are what we differentiate, we do *not*
# return either of their values (since they aren't terribly meaningful,
# thanks to being simultaneously minimized over the model parameters and
# maximized over the internal state). Instead, we just return the value
# of the objective.
output = objective
wrt_objective = tf.cast(distribution[0],
dtype=objective.dtype.base_dtype)
wrt_proxy_constraints = tf.cast(
distribution[1:], dtype=proxy_constraints.dtype.base_dtype)
wrt_state = -self._state_gradient(zero_and_constraints,
distribution)
def gradient_fn(output_gradient):
# We return the gradient w.r.t. the objective, constraints,
# proxy_constraints and internal state, respectively (this is the same
# order as the arguments to loss_gradient_fn). Notice that the gradient
# w.r.t. the constraints is None, and that w.r.t. the internal state is
# scaled by dual_scale.
return (output_gradient * wrt_objective, None,
output_gradient * wrt_proxy_constraints,
self._dual_scale * output_gradient * wrt_state)
return output, gradient_fn
def preprocess_data(image, label, is_training):
"""CIFAR data preprocessing"""
image = tf.image.convert_image_dtype(image, tf.float32)
if is_training:
crop_padding = 4
image = tf.pad(image, [[crop_padding, crop_padding],
[crop_padding, crop_padding], [0, 0]],
'REFLECT')
image = tf.image.random_crop(image, [32, 32, 3])
image = tf.image.random_flip_left_right(image)
if FLAGS.distort_color:
image = color_distortion(image, s=1.0)
else:
image = tf.image.resize_with_crop_or_pad(image, 32, 32) # central crop
return image, label
def resize_and_pad(image, target_size, random_centering):
"""Resize image to target_size (<= image.size) and pad to original size."""
original_shape = image.shape
size = tf.reshape(target_size, [1])
size = tf.concat([size, size], axis=0)
image = tf.image.resize(image, size=size)
pad_size = original_shape[1] - target_size
pad_size_left, pad_size_right = _make_padding_sizes(
pad_size, random_centering)
padding = [[pad_size_left, pad_size_right],
[pad_size_left, pad_size_right], [0, 0]]
if len(original_shape) == 4:
padding = [[0, 0]] + padding
image = tf.pad(image, padding)
image.set_shape(original_shape)
return image
def call(self, inputs):
if self.padding == 'causal':
inputs = tf.pad(inputs, self._compute_causal_padding())
outputs, frequency_l1 = fft_conv1d(inputs, self.kernel)
if self.kernel_regularizer:
self.add_loss(self.kernel_regularizer.l1 * frequency_l1)
if self.use_bias:
outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')
if self.activation is not None:
outputs = self.activation(outputs)
return outputs
def _segmented_range(limits):
"""Equivalent to `tf.ragged.range(limits).flat_values`.
Ragged Tensors are are not supported by numpy.
Args:
limits: Integer `Tensor` of sizes of each range.
Returns:
segments: 1D `Tensor` of segment ranges.
"""
# To cope with [0]-shaped limits, which disagrees with the sensibilities of
# tf.repeat, we left-pad, then slice the output.
limits = tf.pad(limits, [[1, 0]], constant_values=0)
return (tf.range(tf.reduce_sum(limits)) -
tf.repeat(tf.concat([[0], tf.cumsum(limits[:-1])], axis=0), limits))
def _inverse(self, y):
ndims = prefer_static.rank(y)
shifted_y = tf.pad(
tf.slice(
y, tf.zeros(ndims, dtype=tf.int32),
prefer_static.shape(y) -
tf.one_hot(ndims + self.axis, ndims, dtype=tf.int32)
), # Remove the last entry of y in the chosen dimension.
paddings=tf.one_hot(
tf.one_hot(ndims + self.axis, ndims, on_value=0, off_value=-1),
2,
dtype=tf.int32
) # Insert zeros at the beginning of the chosen dimension.
)
return y - shifted_y
def call(self, inputs):
if self.padding == "causal":
inputs = tf.pad(inputs, self._compute_causal_padding(inputs))
if self.data_format == "channels_last":
strides = (1, ) + self.strides * 2 + (1, )
spatial_start_dim = 1
else:
strides = (1, 1) + self.strides * 2
spatial_start_dim = 2
# Explicitly broadcast inputs and kernels to 4D.
# TODO(fchollet): refactor when a native separable_conv1d op is
# available.
inputs = tf.expand_dims(inputs, spatial_start_dim)
depthwise_kernel = tf.expand_dims(self.depthwise_kernel, 0)
pointwise_kernel = tf.expand_dims(self.pointwise_kernel, 0)
dilation_rate = (1, ) + self.dilation_rate
if self.padding == "causal":
op_padding = "valid"
else:
op_padding = self.padding
outputs = tf.compat.v1.nn.separable_conv2d(
inputs,
depthwise_kernel,
pointwise_kernel,
strides=strides,
padding=op_padding.upper(),
rate=dilation_rate,
data_format=conv_utils.convert_data_format(self.data_format,
ndim=4),
)
if self.use_bias:
outputs = tf.nn.bias_add(
outputs,
self.bias,
data_format=conv_utils.convert_data_format(self.data_format,
ndim=4),
)
outputs = tf.squeeze(outputs, [spatial_start_dim])
if self.activation is not None:
return self.activation(outputs)
return outputs
def pad_or_trim_to(x, shape, pad_val=0):
"""Pad and slice x to the given shape.
Args:
x: A tensor.
shape: The shape of the returned tensor.
pad_val: An int or float used to pad x.
Returns:
'x' is padded with pad_val and sliced so that the result has the given
shape.
"""
pad = shape - tf.minimum(tf.shape(x), shape)
zeros = tf.zeros_like(pad)
x = tf.pad(x, tf.stack([zeros, pad], axis=1), constant_values=pad_val)
# If dim-i is larger than shape[i], we slice [0:shape[i]] for dim-i.
return tf.reshape(tf.slice(x, zeros, shape), shape)
def _laplacian_level(self, x, l_w, l_h, l_w2, l_h2):
"""Performs one down/upsampling level of the Laplacian pyramid."""
x_shape = tf.shape(x)
padded = tf.pad(x, ((0, 0), (4, 4), (4, 4), (0, 0)), mode="REFLECT")
down_h = tf.nn.conv2d(
padded, l_h, (2, 1), padding="VALID", data_format="NHWC")
down_hw = tf.nn.conv2d(
down_h, l_w, (1, 2), padding="VALID", data_format="NHWC")
up_h_shape = tf.concat([tf.shape(down_hw)[:-2], x_shape[-2:]], 0)
up_w = tf.nn.conv2d_transpose(
down_hw, l_w2, up_h_shape, (1, 2),
padding=((0, 0), (0, 0), (4, 4), (0, 0)), data_format="NHWC")
up_hw = tf.nn.conv2d_transpose(
up_w, l_h2, x_shape, (2, 1),
padding=((0, 0), (4, 4), (0, 0), (0, 0)), data_format="NHWC")
# For input x^(n), these are called z^(n) and x^(n+1) in the paper.
return x - up_hw, down_hw[:, 1:-1, 1:-1, :]
def pad(array, pad_width, mode, constant_values=0):
"""Pads an array.
Args:
array: array_like of rank N. Input array.
pad_width: {sequence, array_like, int}.
Number of values padded to the edges of each axis.
((before_1, after_1), ... (before_N, after_N)) unique pad widths
for each axis.
((before, after),) yields same before and after pad for each axis.
(pad,) or int is a shortcut for before = after = pad width for all
axes.
mode: string. One of the following string values:
'constant'
Pads with a constant value.
'reflect'
Pads with the reflection of the vector mirrored on
the first and last values of the vector along each
axis.
'symmetric'
Pads with the reflection of the vector mirrored
along the edge of the array.
**NOTE**: The supported list of `mode` does not match that of numpy's.
constant_values: scalar with same dtype as `array`.
Used in 'constant' mode as the pad value. Default is 0.
Returns:
An ndarray padded array of rank equal to `array` with shape increased
according to `pad_width`.
Raises:
ValueError if `mode` is not supported.
"""
if not (mode == 'constant' or mode == 'reflect' or mode == 'symmetric'):
raise ValueError('Unsupported padding mode: ' + mode)
mode = mode.upper()
array = array_creation.asarray(array)
pad_width = array_creation.asarray(pad_width, dtype=tf.int32)
return utils.tensor_to_ndarray(
tf.pad(tensor=array.data,
paddings=pad_width.data,
mode=mode,
constant_values=constant_values))
