def __init__(
self,
mode,
img_shape: Tuple[int, int] = (160, 160),
alpha: float = 0.8,
stride: int = 1,
min_face_size: int = 40,
) -> None:
assert mode in ("mtcnn", "trt-mtcnn"), f"{mode} not supported"
self.mode = mode
self.alpha = alpha
self.img_shape = tuple(img_shape)
self.stride = stride
self.min_face_size = min_face_size
self.frame_ct = 0
self._cached_result = None
if "trt-mtcnn" in mode:
sys.path.insert(1, "../util/trt_mtcnn_plugin")
from util.trt_mtcnn_plugin.trt_mtcnn import TrtMTCNNWrapper # noqa
engine_paths = [
f"../util/trt_mtcnn_plugin/mtcnn/det{i+1}.engine"
for i in range(3)
]
self.trt_mtcnn = TrtMTCNNWrapper(*engine_paths)
if "mtcnn" in mode.replace("trt-mtcnn", ""):
import tensorflow.compat.v1 as tf # noqa
assert tf.executing_eagerly(), (
"[internal] launch failed, tf not eager."
"Check that tensorflow>=2.3 and that eager exec enabled")
mpath = CONFIG_HOME + "/models/mtcnn.pb"
self.mtcnn = tf.wrap_function(
get_mtcnn(mpath, min_size=float(self.min_face_size)),
[tf.TensorSpec(shape=[None, None, 3], dtype=tf.float32)],
)
def call(self, inputs, training):
"""Pass a tensor through the bottleneck.
Args:
inputs: The tensor to be passed through the bottleneck.
training: Boolean. If `True`, returns a differentiable approximation of
the inputs, and their likelihoods under the modeled probability
densities. If `False`, returns the quantized inputs and their
likelihoods under the corresponding probability mass function. These
quantities can't be used for training, as they are not differentiable,
but represent actual compression more closely.
Returns:
values: `Tensor` with the same shape as `inputs` containing the perturbed
or quantized input values.
likelihood: `Tensor` with the same shape as `inputs` containing the
likelihood of `values` under the modeled probability distributions.
Raises:
ValueError: if `inputs` has an integral or inconsistent `DType`, or
inconsistent number of channels.
"""
inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
if inputs.dtype.is_integer:
raise ValueError(
"{} can't take integer inputs.".format(type(self).__name__))
outputs = self._quantize(inputs, "noise" if training else "dequantize")
assert outputs.dtype == self.dtype
likelihood = self._likelihood(outputs)
if self.likelihood_bound > 0:
likelihood_bound = tf.constant(self.likelihood_bound, dtype=self.dtype)
likelihood = math_ops.lower_bound(likelihood, likelihood_bound)
if not tf.executing_eagerly():
outputs_shape, likelihood_shape = self.compute_output_shape(inputs.shape)
outputs.set_shape(outputs_shape)
likelihood.set_shape(likelihood_shape)
return outputs, likelihood
def __init__(self,
population_size,
selection_probability):
"""Initialize the PrivacyLedger.
Args:
population_size: An integer (may be variable) specifying the size of the
population, i.e. size of the training data used in each epoch.
selection_probability: A float (may be variable) specifying the
probability each record is included in a sample.
Raises:
ValueError: If selection_probability is 0.
"""
self._population_size = population_size
self._selection_probability = selection_probability
if tf.executing_eagerly():
if tf.equal(selection_probability, 0):
raise ValueError('Selection probability cannot be 0.')
init_capacity = tf.cast(tf.math.ceil(1 / selection_probability), tf.int32)
else:
if selection_probability == 0:
raise ValueError('Selection probability cannot be 0.')
init_capacity = np.int(np.ceil(1 / selection_probability))
# The query buffer stores rows corresponding to GaussianSumQueryEntries.
self._query_buffer = tensor_buffer.TensorBuffer(
init_capacity, [3], tf.float32, 'query')
self._sample_var = tf.Variable(
initial_value=tf.zeros([3]), trainable=False, name='sample')
# The sample buffer stores rows corresponding to SampleEntries.
self._sample_buffer = tensor_buffer.TensorBuffer(
init_capacity, [3], tf.float32, 'sample')
self._sample_count = tf.Variable(
initial_value=0.0, trainable=False, name='sample_count')
self._query_count = tf.Variable(
initial_value=0.0, trainable=False, name='query_count')
self._cs = tf.CriticalSection()
def _get_matching_collections(
self, mode, tensor, tensor_type, ts_name, is_input_to_model=False, is_output_of_model=False
):
colls_with_tensor = set()
if tensor_type == "weight":
if match_inc(
tensor.name, self.collection_manager.get(CollectionKeys.BIASES).include_regex
):
colls_with_tensor.add(self.collection_manager.get(CollectionKeys.BIASES))
else:
colls_with_tensor.add(self.collection_manager.get(CollectionKeys.WEIGHTS))
elif is_input_to_model:
colls_with_tensor.add(self.collection_manager.get(CollectionKeys.INPUTS))
elif is_output_of_model:
colls_with_tensor.add(self.collection_manager.get(CollectionKeys.OUTPUTS))
for current_coll in self.collection_manager.get_collections().values():
if current_coll.name in [CollectionKeys.WEIGHTS, CollectionKeys.BIASES]:
# don't match regex for these as these are added specially above
# we also don't want users to make mistakes configuring these collections
continue
if match_inc(ts_name, current_coll.include_regex):
# In TF 2.x eager mode, we can't put tensors in a set/dictionary as tensor.__hash__()
# is no longer available. tensor.experimental_ref() returns a hashable reference
# object to this Tensor.
if is_tf_version_2x() and tf.executing_eagerly():
# tensor.experimental_ref is an experimental API
# and can be changed or removed.
# Ref: https://www.tensorflow.org/api_docs/python/tf/Tensor#experimental_ref
tensor = tensor.experimental_ref()
if not current_coll.has_tensor(tensor):
# tensor will be added to this coll below
colls_with_tensor.add(current_coll)
# don't recommend adding tensors externally as
# they will have different internal name
# but regardless, in such case we only use that tensor name to save data
# instead of the keras-style-internal-names
return colls_with_tensor
def _get_weights(model_hparams, vocab_size, hidden_dim=None):
"""Copied from tensor2tensor/layers/modalities.py but uses total vocab."""
if hidden_dim is None:
hidden_dim = model_hparams.hidden_size
num_shards = model_hparams.symbol_modality_num_shards
shards = []
for i in range(num_shards):
shard_size = (sum(vocab_size) // num_shards) + (
1 if i < sum(vocab_size) % num_shards else 0)
var_name = 'weights_%d' % i
shards.append(
tf.get_variable(var_name, [shard_size, hidden_dim],
initializer=tf.random_normal_initializer(
0.0, hidden_dim**-0.5)))
if num_shards == 1:
ret = shards[0]
else:
ret = tf.concat(shards, 0)
# Convert ret to tensor.
if not tf.executing_eagerly():
ret = common_layers.convert_gradient_to_tensor(ret)
return ret
def test_run_in_graph_and_eager_modes_setup_in_same_mode(self):
modes = []
mode_name = lambda: "eager" if tf.executing_eagerly() else "graph"
class ExampleTest(tf.test.TestCase):
def runTest(self):
pass
def setUp(self):
modes.append("setup_" + mode_name())
@test_utils.run_in_graph_and_eager_modes
def testBody(self):
modes.append("run_" + mode_name())
e = ExampleTest()
e.setUp()
e.testBody()
self.assertEqual(modes[0:2], ["setup_eager", "run_eager"])
self.assertEqual(modes[2:], ["setup_graph", "run_graph"])
def test_define_train_ops(self):
if tf.executing_eagerly():
# `tfgan.stargan_model` doesn't work when executing eagerly.
return
hparams = self.hparams._replace(batch_size=2,
generator_lr=0.1,
discriminator_lr=0.01)
images_shape = [hparams.batch_size, 4, 4, 3]
images = tf.zeros(images_shape, dtype=tf.float32)
labels = tf.one_hot([0] * hparams.batch_size, 2)
model = train_lib._define_model(images, labels)
loss = tfgan.stargan_loss(model)
train_ops = train_lib._define_train_ops(model, loss,
hparams.generator_lr,
hparams.discriminator_lr,
hparams.adam_beta1,
hparams.adam_beta2,
hparams.max_number_of_steps)
self.assertIsInstance(train_ops, tfgan.GANTrainOps)
def multiply_gradients_matching_regex(grads_and_vars, regex_list, multiplier):
"""Multiply gradients whose variable names match a regular expression.
Args:
grads_and_vars: A list of gradient to variable pairs (tuples).
regex_list: A list of string regular expressions.
multiplier: A (float) multiplier to apply to each gradient matching the
regular expression.
Returns:
grads_and_vars: A list of gradient to variable pairs (tuples).
"""
if tf.executing_eagerly():
raise ValueError('Accessing variables is not supported in eager mode.')
variables = [pair[1] for pair in grads_and_vars]
matching_vars = filter_variables(variables, regex_list, invert=True)
for var in matching_vars:
logging.info('Applying multiplier %f to variable [%s]',
multiplier, var.op.name)
grad_multipliers = {var: float(multiplier) for var in matching_vars}
return slim.learning.multiply_gradients(grads_and_vars,
grad_multipliers)
def _add_to_device_map(self, tensor):
tensors = []
# In TF 2.x eager mode, we cannot rely on input tensors to
# populate this device map as these tensors cannot be saved.
# Due to this, while executing MirroredStrategy on multiple GPUs,
# weights and biases in the form of values.MirroredVariable are the
# first tensors to reach this point. Since MirroredVariable is not
# processed here, MirroredStrategy distributed training jobs failed
# on GPU. Adding a check and processing MirroredVariable for TF 2.x
# eager mode alone.
if is_tf_version_2x() and tf.executing_eagerly():
from tensorflow.python.distribute import values
if isinstance(tensor, values.DistributedValues):
tensors = [t for t in tensor._values]
else:
tensors = [tensor]
for t in tensors:
if t.device and "CPU" not in t.device and t.device not in self.device_map:
self.device_map[t.device] = serialize_tf_device(t.device)
def call(self, inputs):
inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
ndim = self._input_rank
if self.rectify:
inputs = tf.nn.relu(inputs)
# Compute normalization pool.
if ndim == 2:
norm_pool = tf.linalg.matmul(tf.math.square(inputs), self.gamma)
norm_pool = tf.nn.bias_add(norm_pool, self.beta)
elif self.data_format == "channels_last" and ndim <= 4:
# TODO(unassigned): This branch should also work for ndim == 5, but
# currently triggers a bug in TF.
shape = self.gamma.shape.as_list()
gamma = tf.reshape(self.gamma, (ndim - 2) * [1] + shape)
norm_pool = tf.nn.convolution(tf.math.square(inputs), gamma,
"VALID")
norm_pool = tf.nn.bias_add(norm_pool, self.beta)
else: # generic implementation
# This puts channels in the last dimension regardless of input.
norm_pool = tf.linalg.tensordot(tf.math.square(inputs), self.gamma,
[[self._channel_axis()], [0]])
norm_pool += self.beta
if self.data_format == "channels_first":
# Return to channels_first format if necessary.
axes = list(range(ndim - 1))
axes.insert(1, ndim - 1)
norm_pool = tf.transpose(norm_pool, axes)
if self.inverse:
norm_pool = tf.math.sqrt(norm_pool)
else:
norm_pool = tf.math.rsqrt(norm_pool)
outputs = inputs * norm_pool
if not tf.executing_eagerly():
outputs.set_shape(self.compute_output_shape(inputs.shape))
return outputs
def test_generator_shapes_and_ranges(self):
"""Tests the discriminator.
Make sure the image shapes and output value ranges are as expected.
"""
if tf.executing_eagerly():
# `compute_spectral_norm` doesn't work when executing eagerly.
return
batch_size = 10
num_classes = 1000
gen_class_logits = tf.zeros((batch_size, num_classes))
gen_class_ints = tf.random.categorical(logits=gen_class_logits,
num_samples=1)
gen_sparse_class = tf.squeeze(gen_class_ints)
images = tf.random.normal([10, 32, 32, 3])
d_out, var_list = discriminator.discriminator(images, gen_sparse_class,
16, 1000)
sess = tf.train.MonitoredTrainingSession()
images_np = sess.run(d_out)
self.assertEqual((batch_size, 1), images_np.shape)
self.assertAllInRange(images_np, -1.0, 1.0)
self.assertIsInstance(var_list, list)
def get_shape_list(tensor, expected_rank=None, name=None):
"""Returns a list of the shape of tensor, preferring static dimensions.
Args:
tensor: A tf.Tensor object to find the shape of.
expected_rank: (optional) int. The expected rank of `tensor`. If this is
specified and the `tensor` has a different rank, and exception will be
thrown.
name: Optional name of the tensor for the error message.
Returns:
A list of dimensions of the shape of tensor. All static dimensions will
be returned as python integers, and dynamic dimensions will be returned
as tf.Tensor scalars.
"""
if name is None:
# Tensor.name is not supported in Eager mode.
if tf.executing_eagerly():
name = "get_shape_list"
else:
name = tensor.name
if expected_rank is not None:
assert_rank(tensor, expected_rank, name)
shape = tensor.shape.as_list()
non_static_indexes = []
for (index, dim) in enumerate(shape):
if dim is None:
non_static_indexes.append(index)
if not non_static_indexes:
return shape
dyn_shape = tf.shape(tensor)
for index in non_static_indexes:
shape[index] = dyn_shape[index]
return shape
def testPartitioners(self):
if tf.executing_eagerly():
self.skipTest("Eager does not support partitioned variables.")
partitioners = {
"w": tf.fixed_size_partitioner(num_shards=2),
"b": tf.fixed_size_partitioner(num_shards=2),
}
alex_net = snt.nets.AlexNetMini(
partitioners=partitioners, name="alexnet1")
input_shape = [alex_net._min_size, alex_net._min_size, 3]
inputs = tf.placeholder(tf.float32, shape=[None] + input_shape)
alex_net(inputs)
for conv_module in alex_net.conv_modules:
self.assertEqual(type(conv_module.w), variables.PartitionedVariable)
self.assertEqual(type(conv_module.b), variables.PartitionedVariable)
for linear_module in alex_net.linear_modules:
self.assertEqual(type(linear_module.w), variables.PartitionedVariable)
self.assertEqual(type(linear_module.b), variables.PartitionedVariable)
def testGetAllVariablesWithConditionalConstruction(self):
inputs = tf.ones(dtype=tf.float32, shape=[10, 10])
cond = tf.constant(0.)
module_a = SimpleModule(name="module_a")
module_b = SimpleModule(name="module_b")
_ = tf.cond(cond > 0, lambda: module_a(inputs), lambda: module_b(inputs)) # pylint: disable=not-callable
if tf.executing_eagerly():
# In eager mode only the true branch is taken.
msg = "module_a not instantiated yet"
with self.assertRaisesRegexp(base.NotConnectedError, msg):
module_a.get_all_variables()
else:
# check module_a
all_variables = module_a.get_all_variables()
all_variable_names = sorted([str(v.name) for v in all_variables])
self.assertEqual(["module_a/b:0", "module_a/w:0"], all_variable_names)
# check module_b
all_variables = module_b.get_all_variables()
all_variable_names = sorted([str(v.name) for v in all_variables])
self.assertEqual(["module_b/b:0", "module_b/w:0"], all_variable_names)
def testRegularizersInRegularizationLosses(self):
w_regularizer = contrib_layers.l1_regularizer(scale=0.5)
b_regularizer = contrib_layers.l2_regularizer(scale=0.5)
self.setUpWithNumOutputClasses(1)
dilation_mod = snt.nets.Dilation(
num_output_classes=self._num_output_classes,
regularizers={
"w": w_regularizer,
"b": b_regularizer
})
dilation_mod(tf.convert_to_tensor(self._images))
regularizers = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
# There are two regularizers per level
layers_number = 8
self.assertEqual(len(regularizers), 2 * layers_number)
if not tf.executing_eagerly():
for i in range(0, 2 * layers_number, 2):
self.assertRegexpMatches(regularizers[i].name,
".*l1_regularizer.*")
self.assertRegexpMatches(regularizers[i + 1].name,
".*l2_regularizer.*")
def test_build_graph(self, mock_data_provider):
if tf.executing_eagerly():
# `tfgan.gan_model` doesn't work when executing eagerly.
return
hparams = train_lib.HParams(
batch_size=16,
max_number_of_steps=0,
generator_lr=0.0002,
discriminator_lr=0.0002,
master='',
train_log_dir='/tmp/tfgan_logdir/cifar/',
ps_replicas=0,
task=0)
# Mock input pipeline.
mock_imgs = np.zeros([hparams.batch_size, 32, 32, 3], dtype=np.float32)
mock_lbls = np.concatenate(
(np.ones([hparams.batch_size, 1], dtype=np.float32),
np.zeros([hparams.batch_size, 9], dtype=np.float32)),
axis=1)
mock_data_provider.provide_data.return_value = (mock_imgs, mock_lbls)
train_lib.train(hparams)
def _is_not_supported(self):
if self.distribution_strategy is None:
self.distribution_strategy = self._get_distribution_strategy()
if self._hook_supported is None:
self._hook_supported = True
if tf.executing_eagerly() or (
hasattr(self.model, "run_eagerly") and self.model.run_eagerly
):
if is_tf_version_2x():
self.logger.info(
"Executing in TF2.x eager mode."
"TF 2.x eager doesn't provide gradient and optimizer variable values."
"SageMaker Debugger will not be saving gradients and optimizer variables in this case"
)
else:
self.logger.info(
"Disabling SMDebug as it does not support eager mode" "for TF versions 1.x"
)
self._hook_supported = False
elif self.distribution_strategy == TFDistributionStrategy.MIRRORED:
try:
from tensorflow.python.keras.distribute.distributed_training_utils import (
get_distributed_model,
)
except ImportError:
# for tf1.13 we can't import this, so we can't support mirrored strategy
self.logger.info(
"Disabling SMDebug as it does not support mirrored strategy"
"with TensorFlow version <1.14"
)
self._hook_supported = False
elif self.distribution_strategy == TFDistributionStrategy.UNSUPPORTED:
self.logger.info(
f"Disabling SMDebug as it does not support " f"{tf.distribute.get_strategy()}"
)
self._hook_supported = False
return not self._hook_supported
def __init__(self,
tail_mass=2**-8,
likelihood_bound=1e-9,
range_coder_precision=16,
**kwargs):
"""Initializer.
Arguments:
tail_mass: Float, between 0 and 1. The bottleneck layer automatically
determines the range of input values based on their frequency of
occurrence. Values occurring in the tails of the distributions will not
be encoded with range coding, but using a Golomb-like code. `tail_mass`
determines the amount of probability mass in the tails which will be
Golomb-coded. For example, the default value of `2 ** -8` means that on
average, one 256th of all values will use the Golomb code.
likelihood_bound: Float. If positive, the returned likelihood values are
ensured to be greater than or equal to this value. This prevents very
large gradients with a typical entropy loss (defaults to 1e-9).
range_coder_precision: Integer, between 1 and 16. The precision of the
range coder used for compression and decompression. This trades off
computation speed with compression efficiency, where 16 is the slowest
but most efficient setting. Choosing lower values may increase the
average codelength slightly compared to the estimated entropies.
**kwargs: Other keyword arguments passed to superclass (`Layer`).
"""
super(EntropyModel, self).__init__(**kwargs)
self._tail_mass = float(tail_mass)
if not 0 < self.tail_mass < 1:
raise ValueError(
"`tail_mass` must be between 0 and 1, got {}.".format(
self.tail_mass))
self._likelihood_bound = float(likelihood_bound)
self._range_coder_precision = int(range_coder_precision)
if tf.executing_eagerly():
raise NotImplementedError(
"Keras layer implementations of entropy models are not supported in "
"eager mode.")
def get_variables(self,
collection=tf.GraphKeys.TRAINABLE_VARIABLES):
"""Returns tuple of `tf.Variable`s declared inside this module.
Note that this operates by searching this module's variable scope,
and so does not know about any modules that were constructed elsewhere but
used inside this module.
This method explicitly re-enters the Graph which this module has been
connected to.
Args:
collection: Collection to restrict query to. By default this is
`tf.GraphKeys.TRAINABLE_VARIABLES`, which doesn't
include non-trainable variables such as moving averages.
Returns:
A tuple of `tf.Variable` objects.
Raises:
NotConnectedError: If the module is not connected to the Graph.
"""
self._ensure_is_connected()
if self._defun_wrapped and tf.executing_eagerly():
raise NotSupportedError(
"`module.get_variables()` relies on TensorFlow collections which are "
"not supported when your module is wrapped with defun. Instead use "
"`module.trainable_variables` or `module.variables`.")
# Explicitly re-enter Graph, in case the module is being queried with a
# different default Graph from the one it was connected to. If this was not
# here then querying the variables from a different graph scope would
# produce an empty tuple.
with self._graph.as_default():
return util.get_variables_in_scope(
self.variable_scope, collection=collection)
def testInitialStateNames(self):
if tf.executing_eagerly():
return self.skipTest("Tensor.name is meaningless in eager mode.")
hidden_size_a = 3
hidden_size_b = 4
batch_size = 5
deep_rnn = snt.DeepRNN([
snt.LSTM(hidden_size_a, name="a"),
snt.LSTM(hidden_size_b, name="b")
])
deep_rnn_state = deep_rnn.initial_state(batch_size, trainable=True)
self.assertEqual(
deep_rnn_state[0][0].name,
"deep_rnn_initial_state/a_initial_state/state_hidden_tiled:0")
self.assertEqual(
deep_rnn_state[0][1].name,
"deep_rnn_initial_state/a_initial_state/state_cell_tiled:0")
self.assertEqual(
deep_rnn_state[1][0].name,
"deep_rnn_initial_state/b_initial_state/state_hidden_tiled:0")
self.assertEqual(
deep_rnn_state[1][1].name,
"deep_rnn_initial_state/b_initial_state/state_cell_tiled:0")
other_start_state = deep_rnn.initial_state(batch_size,
trainable=True,
name="blah")
self.assertEqual(other_start_state[0][0].name,
"blah/a_initial_state/state_hidden_tiled:0")
self.assertEqual(other_start_state[0][1].name,
"blah/a_initial_state/state_cell_tiled:0")
self.assertEqual(other_start_state[1][0].name,
"blah/b_initial_state/state_hidden_tiled:0")
self.assertEqual(other_start_state[1][1].name,
"blah/b_initial_state/state_cell_tiled:0")
def testSubgraphsNotRecordedEager(self):
if not tf.executing_eagerly():
self.skipTest("Subgraphs are recorded in graph mode")
id_mod = IdentityModule(name="foo")
with self.assertRaisesRegexp(base.NotSupportedError,
"not tracked in eager mode"):
id_mod.last_connected_subgraph()
# pylint: disable=not-callable
inputs = tf.ones(dtype=tf.float32, shape=[21])
id_mod(inputs)
with tf.name_scope("blah"):
blah_inputs = tf.ones(dtype=tf.float32, shape=[22])
id_mod(blah_inputs)
with tf.name_scope("baz"):
baz_inputs = tf.ones(dtype=tf.float32, shape=[23])
id_mod(baz_inputs)
# pylint: enable=not-callable
with self.assertRaisesRegexp(base.NotSupportedError,
"not tracked in eager mode"):
id_mod.connected_subgraphs # pylint: disable=pointless-statement
def get_variable_initializer(hparams):
"""Get variable initializer from hparams."""
if not hparams.initializer:
return None
if not tf.executing_eagerly():
tf.logging.info("Using variable initializer: %s", hparams.initializer)
if hparams.initializer == "orthogonal":
return tf.orthogonal_initializer(gain=hparams.initializer_gain)
elif hparams.initializer == "uniform":
max_val = 0.1 * hparams.initializer_gain
return tf.random_uniform_initializer(-max_val, max_val)
elif hparams.initializer == "normal_unit_scaling":
return tf.variance_scaling_initializer(hparams.initializer_gain,
mode="fan_avg",
distribution="normal")
elif hparams.initializer == "uniform_unit_scaling":
return tf.variance_scaling_initializer(hparams.initializer_gain,
mode="fan_avg",
distribution="uniform")
elif hparams.initializer == "xavier":
return tf.initializers.glorot_uniform()
else:
raise ValueError("Unrecognized initializer: %s" % hparams.initializer)
def inference_network_fn(self,
features,
labels,
mode,
config=None,
params=None):
"""See base class documentation."""
del mode, config, params
if not self._model:
self._build_model()
if self._multi_dataset:
if tf.executing_eagerly():
x1 = tf.convert_to_tensor(features.x1)
x2 = tf.convert_to_tensor(features.x2)
else:
x1 = features.x1
x2 = features.x2
net = x1 + x2
else:
net = features.x
net = self._model(net)
return dict(logits=net)
def testPartitioners(self):
if tf.executing_eagerly():
self.skipTest("Eager does not support partitioned variables.")
partitioners = {
"w": tf.variable_axis_size_partitioner(10),
"b": tf.variable_axis_size_partitioner(8),
}
module = snt.nets.ConvNet2DTranspose(output_channels=self.output_channels,
output_shapes=self.output_shapes,
kernel_shapes=self.kernel_shapes,
strides=self.strides,
paddings=self.paddings,
partitioners=partitioners)
input_shape = [10, 100, 100, 3]
input_to_net = tf.placeholder(tf.float32, shape=input_shape)
_ = module(input_to_net)
for layer in module._layers:
self.assertEqual(type(layer.w), variables.PartitionedVariable)
self.assertEqual(type(layer.b), variables.PartitionedVariable)
def testLearnableMultivariateNormalDiagCellClassNoBatch(self):
prior = disentangled_vae.LearnableMultivariateNormalDiagCell(
self.dimensions, self.hidden_size)
# Zero state.
dynamic_previous_output, state = prior.zero_state()
self.assertEqual(dynamic_previous_output.shape, (self.dimensions, ))
for tensor in state:
self.assertEqual(tensor.shape, (1, self.hidden_size))
h0, c0 = state
# First timestep.
dist_z1, state_z1 = prior(dynamic_previous_output, state)
self.assertDistShape(dist_z1, (self.dimensions, ), ())
for tensor in state_z1:
self.assertEqual(tensor.shape, (1, self.hidden_size))
if not tf.executing_eagerly():
self.evaluate(tf.compat.v1.global_variables_initializer())
h1, c1 = state_z1
self.assertTrue(np.allclose(self.evaluate(h1), self.evaluate(h0)))
self.assertTrue(np.allclose(self.evaluate(c1), self.evaluate(c0)))
# Second timestep.
dist_z2, state = prior(dist_z1.sample(), state_z1)
self.assertDistShape(dist_z2, (self.dimensions, ), ())
for tensor in state:
self.assertEqual(tensor.shape, (1, self.hidden_size))
h2, c2 = state
self.assertFalse(np.allclose(self.evaluate(h2), self.evaluate(h1)))
self.assertFalse(np.allclose(self.evaluate(c2), self.evaluate(c1)))
# Second timestep with sample shape.
dist_z2, state = prior(dist_z1.sample(2), state_z1)
self.assertDistShape(dist_z2, (self.dimensions, ), (2))
for tensor in state:
self.assertEqual(tensor.shape, (2, self.hidden_size))
def _slow_greedy_infer(self, features, decode_length):
"""A slow greedy inference method.
Quadratic time in decode_length.
Args:
features: an map of string to `Tensor`
decode_length: an integer. How many additional timesteps to decode.
Returns:
A dict of decoding results {
"outputs": integer `Tensor` of decoded ids of shape
[batch_size, <= decode_length] if beam_size == 1 or
[batch_size, top_beams, <= decode_length]
"scores": None
"logits": `Tensor` of shape [batch_size, time, 1, 1, vocab_size].
"losses": a dictionary: {loss-name (string): floating point `Scalar`}
}
"""
if not features:
features = {}
inputs_old = None
# process all conditioning features
if "inputs" in features:
if len(features["inputs"].shape) < 4:
inputs_old = features["inputs"]
features["inputs"] = tf.expand_dims(features["inputs"], 2)
else: # this would be for melody decoding
if "melody" in features:
if len(features["melody"].shape) < 4:
inputs_old = features["melody"]
features["melody"] = tf.expand_dims(features["melody"], 2)
if "performance" in features:
if len(features["performance"].shape) < 4:
inputs_old = features["performance"]
features["performance"] = tf.expand_dims(
features["performance"], 2)
if not self.has_input:
# Prepare partial targets.
# In either features["inputs"] or features["targets"].
# We force the outputs to begin with these sequences.
partial_targets = features.get("inputs")
if partial_targets is None:
partial_targets = features["targets"]
features["partial_targets"] = tf.to_int64(partial_targets)
# Save the targets in a var and reassign it after the tf.while loop to avoid
# having targets being in a 'while' frame. This ensures targets when used
# in metric functions stays in the same frame as other vars.
targets_old = features.get("targets", None)
target_modality = self._problem_hparams.modality["targets"]
def infer_step(recent_output, recent_logits, unused_loss):
"""Inference step."""
if not tf.executing_eagerly():
if self._target_modality_is_real:
dim = self._problem_hparams.vocab_size["targets"]
if dim is not None and hasattr(self._hparams,
"vocab_divisor"):
dim += (-dim) % self._hparams.vocab_divisor
recent_output.set_shape([None, None, None, dim])
else:
recent_output.set_shape([None, None, None, 1])
padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
features["targets"] = padded
# This is inefficient in that it generates samples at all timesteps,
# not just the last one, except if target_modality is pointwise.
samples, logits, losses = self.sample(features)
# Concatenate the already-generated recent_output with last timestep
# of the newly-generated samples.
top = self._hparams.top.get("targets",
modalities.get_top(target_modality))
if getattr(top, "pointwise", False):
cur_sample = samples[:, -1, :, :]
else:
cur_sample = samples[:,
common_layers.shape_list(recent_output
)[1], :, :]
if self._target_modality_is_real:
cur_sample = tf.expand_dims(cur_sample, axis=1)
samples = tf.concat([recent_output, cur_sample], axis=1)
else:
cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
samples = tf.concat([recent_output, cur_sample], axis=1)
if not tf.executing_eagerly():
samples.set_shape([None, None, None, 1])
# Assuming we have one shard for logits.
logits = tf.concat([recent_logits, logits[:, -1:]], 1)
loss = sum([l for l in losses.values() if l is not None])
return samples, logits, loss
# Create an initial output tensor. This will be passed
# to the infer_step, which adds one timestep at every iteration.
if "partial_targets" in features:
initial_output = tf.to_int64(features["partial_targets"])
while len(initial_output.get_shape().as_list()) < 4:
initial_output = tf.expand_dims(initial_output, 2)
batch_size = common_layers.shape_list(initial_output)[0]
else:
batch_size = common_layers.shape_list(features["performance"])[0]
if self._target_modality_is_real:
dim = self._problem_hparams.vocab_size["targets"]
if dim is not None and hasattr(self._hparams, "vocab_divisor"):
dim += (-dim) % self._hparams.vocab_divisor
initial_output = tf.zeros((batch_size, 0, 1, dim),
dtype=tf.float32)
else:
initial_output = tf.zeros((batch_size, 0, 1, 1),
dtype=tf.int64)
# Hack: foldl complains when the output shape is less specified than the
# input shape, so we confuse it about the input shape.
initial_output = tf.slice(initial_output, [0, 0, 0, 0],
common_layers.shape_list(initial_output))
target_modality = self._problem_hparams.modality["targets"]
if target_modality == modalities.ModalityType.CLASS_LABEL:
decode_length = 1
else:
if "partial_targets" in features:
prefix_length = common_layers.shape_list(
features["partial_targets"])[1]
else:
# this code will generate outputs that tend to be long,
# but this is to avoid the case when the melody is extremely short.
# this can be changed to features["melody"] for the actual behavior.
prefix_length = common_layers.shape_list(
features["performance"])[1]
decode_length = prefix_length + decode_length
# Initial values of result, logits and loss.
result = initial_output
vocab_size = self._problem_hparams.vocab_size["targets"]
if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
vocab_size += (-vocab_size) % self._hparams.vocab_divisor
if self._target_modality_is_real:
logits = tf.zeros((batch_size, 0, 1, vocab_size))
logits_shape_inv = [None, None, None, None]
else:
# tensor of shape [batch_size, time, 1, 1, vocab_size]
logits = tf.zeros((batch_size, 0, 1, 1, vocab_size))
logits_shape_inv = [None, None, None, None, None]
if not tf.executing_eagerly():
logits.set_shape(logits_shape_inv)
loss = 0.0
def while_exit_cond(result, logits, loss): # pylint: disable=unused-argument
"""Exit the loop either if reach decode_length or EOS."""
length = common_layers.shape_list(result)[1]
not_overflow = length < decode_length
if self._problem_hparams.stop_at_eos:
def fn_not_eos():
return tf.not_equal( # Check if the last predicted element is a EOS
tf.squeeze(result[:, -1, :, :]), text_encoder.EOS_ID)
not_eos = tf.cond(
# We only check for early stopping if there is at least 1 element (
# otherwise not_eos will crash).
tf.not_equal(length, 0),
fn_not_eos,
lambda: True,
)
return tf.cond(
tf.equal(batch_size, 1),
# If batch_size == 1, we check EOS for early stopping.
lambda: tf.logical_and(not_overflow, not_eos),
# Else, just wait for max length
lambda: not_overflow)
return not_overflow
result, logits, loss = tf.while_loop(
while_exit_cond,
infer_step, [result, logits, loss],
shape_invariants=[
tf.TensorShape([None, None, None, None]),
tf.TensorShape(logits_shape_inv),
tf.TensorShape([]),
],
back_prop=False,
parallel_iterations=1)
if inputs_old is not None: # Restore to not confuse Estimator.
features["inputs"] = inputs_old
# Reassign targets back to the previous value.
if targets_old is not None:
features["targets"] = targets_old
losses = {"training": loss}
if "partial_targets" in features:
partial_target_length = common_layers.shape_list(
features["partial_targets"])[1]
result = tf.slice(result, [0, partial_target_length, 0, 0],
[-1, -1, -1, -1])
return {
"outputs": result,
"scores": None,
"logits": logits,
"losses": losses,
}
def _learning_rate_return_value(eager_decay_rate):
"""Helper function to return proper learning rate based on tf version."""
if tf.executing_eagerly():
return eager_decay_rate
else:
return eager_decay_rate()
import tvm
from tvm import te
from tvm import relay
from tvm.contrib import graph_runtime
from tvm.relay.testing.config import ctx_list
import keras
try:
import tensorflow.compat.v1 as tf
except ImportError:
import tensorflow as tf
from tensorflow import keras as tf_keras
from packaging import version as package_version
# prevent Keras from using up all gpu memory
if tf.executing_eagerly():
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
else:
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
set_session(tf.Session(config=config))
def pytest_generate_tests(metafunc):
# This function generates the list of tests for pytest, based
# on scenatios that will change the parameters in which the
# tests use to run.
# https://docs.pytest.org/en/latest/example/parametrize.html
def cosine_decay_with_warmup(global_step,
learning_rate_base,
total_steps,
warmup_learning_rate=0.0,
warmup_steps=0,
hold_base_rate_steps=0):
"""Cosine decay schedule with warm up period.
Cosine annealing learning rate as described in:
Loshchilov and Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts.
ICLR 2017. https://arxiv.org/abs/1608.03983
In this schedule, the learning rate grows linearly from warmup_learning_rate
to learning_rate_base for warmup_steps, then transitions to a cosine decay
schedule.
Args:
global_step: int64 (scalar) tensor representing global step.
learning_rate_base: base learning rate.
total_steps: total number of training steps.
warmup_learning_rate: initial learning rate for warm up.
warmup_steps: number of warmup steps.
hold_base_rate_steps: Optional number of steps to hold base learning rate
before decaying.
Returns:
If executing eagerly:
returns a no-arg callable that outputs the (scalar)
float tensor learning rate given the current value of global_step.
If in a graph:
immediately returns a (scalar) float tensor representing learning rate.
Raises:
ValueError: if warmup_learning_rate is larger than learning_rate_base,
or if warmup_steps is larger than total_steps.
"""
if total_steps < warmup_steps:
raise ValueError(
'total_steps must be larger or equal to warmup_steps.')
def eager_decay_rate():
"""Callable to compute the learning rate."""
learning_rate = 0.5 * learning_rate_base * (1 + tf.cos(
np.pi * (tf.cast(global_step, tf.float32) - warmup_steps -
hold_base_rate_steps) /
float(total_steps - warmup_steps - hold_base_rate_steps)))
if hold_base_rate_steps > 0:
learning_rate = tf.where(
global_step > warmup_steps + hold_base_rate_steps,
learning_rate, learning_rate_base)
if warmup_steps > 0:
if learning_rate_base < warmup_learning_rate:
raise ValueError(
'learning_rate_base must be larger or equal to warmup_learning_rate.'
)
slope = (learning_rate_base - warmup_learning_rate) / warmup_steps
warmup_rate = slope * tf.cast(global_step,
tf.float32) + warmup_learning_rate
learning_rate = tf.where(global_step < warmup_steps, warmup_rate,
learning_rate)
return tf.where(global_step > total_steps,
0.0,
learning_rate,
name='learning_rate')
if tf.executing_eagerly():
return eager_decay_rate
else:
return eager_decay_rate()
def manual_stepping(global_step, boundaries, rates, warmup=False):
"""Manually stepped learning rate schedule.
This function provides fine grained control over learning rates. One must
specify a sequence of learning rates as well as a set of integer steps
at which the current learning rate must transition to the next. For example,
if boundaries = [5, 10] and rates = [.1, .01, .001], then the learning
rate returned by this function is .1 for global_step=0,...,4, .01 for
global_step=5...9, and .001 for global_step=10 and onward.
Args:
global_step: int64 (scalar) tensor representing global step.
boundaries: a list of global steps at which to switch learning
rates. This list is assumed to consist of increasing positive integers.
rates: a list of (float) learning rates corresponding to intervals between
the boundaries. The length of this list must be exactly
len(boundaries) + 1.
warmup: Whether to linearly interpolate learning rate for steps in
[0, boundaries[0]].
Returns:
If executing eagerly:
returns a no-arg callable that outputs the (scalar)
float tensor learning rate given the current value of global_step.
If in a graph:
immediately returns a (scalar) float tensor representing learning rate.
Raises:
ValueError: if one of the following checks fails:
1. boundaries is a strictly increasing list of positive integers
2. len(rates) == len(boundaries) + 1
3. boundaries[0] != 0
"""
if any([b < 0 for b in boundaries]) or any(
[not isinstance(b, int) for b in boundaries]):
raise ValueError('boundaries must be a list of positive integers')
if any([bnext <= b for bnext, b in zip(boundaries[1:], boundaries[:-1])]):
raise ValueError('Entries in boundaries must be strictly increasing.')
if any([not isinstance(r, float) for r in rates]):
raise ValueError('Learning rates must be floats')
if len(rates) != len(boundaries) + 1:
raise ValueError('Number of provided learning rates must exceed '
'number of boundary points by exactly 1.')
if boundaries and boundaries[0] == 0:
raise ValueError('First step cannot be zero.')
if warmup and boundaries:
slope = (rates[1] - rates[0]) * 1.0 / boundaries[0]
warmup_steps = list(range(boundaries[0]))
warmup_rates = [rates[0] + slope * step for step in warmup_steps]
boundaries = warmup_steps + boundaries
rates = warmup_rates + rates[1:]
else:
boundaries = [0] + boundaries
num_boundaries = len(boundaries)
def eager_decay_rate():
"""Callable to compute the learning rate."""
rate_index = tf.reduce_max(
tf.where(tf.greater_equal(global_step, boundaries),
list(range(num_boundaries)), [0] * num_boundaries))
return tf.reduce_sum(rates *
tf.one_hot(rate_index, depth=num_boundaries),
name='learning_rate')
if tf.executing_eagerly():
return eager_decay_rate
else:
return eager_decay_rate()