def make_const(
g, # type: base_graph.BaseGraph
name, # type: str
value, # type: np.ndarray
uniquify_name=False # type: bool
):
"""
Convenience method to add a `Const` op to a `gde.Graph`.
Args:
g: The graph that the node should be added to
name: Name for the new `Const` node
value: Value to use for the constant
uniquify_name: if True, generate unique names by appending a numeric
suffix in the event of a name collision. Otherwise name collisions
result in an error.
Returns `gde.Node` object representing the new node.
"""
dtype = tf.as_dtype(value.dtype)
ret = g.add_node(name, "Const", uniquify_name=uniquify_name)
ret.add_attr("dtype", dtype)
ret.add_attr("value", value)
ret.set_outputs_from_pairs([(dtype, tf.TensorShape(value.shape))])
return ret
def _legacy_output_transform_func(*expr,
out_mul=1.0,
out_add=0.0,
out_shrink=1,
out_dtype=None):
if out_mul != 1.0:
expr = [x * out_mul for x in expr]
if out_add != 0.0:
expr = [x + out_add for x in expr]
if out_shrink > 1:
ksize = [1, 1, out_shrink, out_shrink]
expr = [
tf.nn.avg_pool(x,
ksize=ksize,
strides=ksize,
padding="VALID",
data_format="NCHW") for x in expr
]
if out_dtype is not None:
if tf.as_dtype(out_dtype).is_integer:
expr = [tf.round(x) for x in expr]
expr = [tf.saturate_cast(x, out_dtype) for x in expr]
return expr
def make_optimizer(lr, last_itr):
with tf.variable_scope("training", reuse=tf.AUTO_REUSE, use_resource=True):
optimizer_class, optimizer_kwargs = build_optimizer(opts.optimizer, opts.optimizer_arg)
optimizer_class = optimizers.SparseOptimizer(optimizer_class)
optimizer_class = global_step_update_opt.GlobalStepUpdateOptimizer(optimizer_class)
if opts.loss_scale != 1:
optimizer_class = scaling_opt.LossScalingOptimizer(optimizer_class)
optimizer_kwargs['loss_scale'] = opts.loss_scale
optimizer_kwargs['unscale_grad_pre_acc'] = opts.unscale_grad_pre_acc
if opts.grad_acculation_mode == 'Avg':
optimizer_class = scaling_opt.GradScalingOptimizer(optimizer_class)
optimizer_kwargs['grad_scale'] = 1 / opts.gradient_accumulation_count
optimizer_kwargs['scale_grad_pre_acc'] = opts.scale_grad_pre_acc
if opts.grad_norm_clip:
optimizer_class = grad_clip_opt.GradientClippingOptimizer(optimizer_class)
optimizer_kwargs['norm_clip_threshold'] = opts.grad_norm_clip
if opts.slots_fp_type is not None and tf.as_dtype(opts.slots_fp_type) != opts.dtype:
optimizer_class = fp_slot_opt.SelectableSlotFPFormatOptimizer(optimizer_class)
optimizer_kwargs['slots_dtype'] = opts.slots_fp_type
optimizer_kwargs['force_fp32_weight_update'] = opts.force_fp32_weight_update
optimizer = optimizer_class(learning_rate=lr, **optimizer_kwargs,
sparse_layers=transformer.sparse_layers.values(),
dense_gradient_condition=enable_dense_grad and last_itr,
prune_and_grow_outfeed=dense_queue)
return optimizer
def reset_fn(**kwargs):
"""A reset_fn for SV2P.
Args:
**kwargs: a dictionary of inputs, including previous state.
Returns:
a new dictionary with posteriors removed from the state.
"""
batch_size = kwargs["proposals"]
input_len = model.input_length
image_shape = model.frame_size
if model.is_discrete_action:
action_shape = (model.action_space.n, )
else:
action_shape = model.action_space.shape
action_dtype = tf.float32
return {
"images":
tf.zeros((batch_size, input_len) + image_shape, dtype=tf.int32),
"actions":
tf.zeros((batch_size, input_len) + action_shape,
dtype=tf.as_dtype(action_dtype)),
"rewards":
tf.zeros((batch_size, input_len, 1)),
}
def gym_space_spec(gym_space):
"""Returns a reading spec of a gym space.
NOTE: Only implemented currently for Box and Discrete.
Args:
gym_space: instance of gym.spaces whose spec we want.
Returns:
Reading spec for that space.
Raises:
NotImplementedError: For spaces whose reading spec we haven't implemented.
"""
# First try to determine the type.
try:
tf_dtype = tf.as_dtype(gym_space.dtype)
except TypeError as e:
tf.logging.error("Cannot convert space's type [%s] to tf.dtype",
gym_space.dtype)
raise e
# Now hand it over to the specialized functions.
if isinstance(gym_space, Box):
return box_space_spec(gym_space, tf_dtype)
elif isinstance(gym_space, Discrete):
return discrete_space_spec(gym_space, tf_dtype)
else:
raise NotImplementedError
def build(self, input_shape):
"""Implements build() for the layer."""
dtype = tf.as_dtype(self.dtype or tf.keras.backend.floatx())
if not (dtype.is_floating or dtype.is_complex):
raise TypeError("Unable to build `Dense2DProjection` layer with "
"non-floating point (and non-complex) "
"dtype %s" % (dtype, ))
input_shape = tf.TensorShape(input_shape)
if tf.compat.dimension_value(input_shape[-1]) is None:
raise ValueError("The last dimension of the inputs to "
"`Dense2DProjection` should be defined. "
"Found `None`.")
last_dim = tf.compat.dimension_value(input_shape[-1])
self.input_spec = tf.keras.layers.InputSpec(min_ndim=3,
axes={-1: last_dim})
self.kernel = self.add_weight("kernel",
shape=[last_dim, self.output_size],
initializer=self.kernel_initializer,
dtype=self.dtype,
trainable=True)
self.bias = self.add_weight("bias",
shape=[self.output_size],
initializer=self.bias_initializer,
dtype=self.dtype,
trainable=True)
super(Dense2DProjection, self).build(input_shape)
def save_numpy_arrays_to_checkpoint(checkpoint_path, **dict_with_arrays):
"""Saves several NumpyArrays to variables in a TF checkpoint.
Args:
checkpoint_path: String with the path to the checkpoint file.
**dict_with_arrays: Dictionary with keys that signify variable names and
values that are the corresponding Numpy arrays to be saved.
"""
with tf.Graph().as_default():
feed_dict = {}
assign_ops = []
nodes_to_save = []
for array_name, array in dict_with_arrays.items():
# We will save the numpy array with the corresponding dtype.
tf_dtype = tf.as_dtype(array.dtype)
# We create a variable which we would like to persist in the checkpoint.
node = tf.get_variable(array_name, shape=array.shape, dtype=tf_dtype)
nodes_to_save.append(node)
# We feed the numpy arrays into the graph via placeholder which avoids
# adding the numpy arrays to the graph as constants.
placeholder = tf.placeholder(tf_dtype, shape=array.shape)
feed_dict[placeholder] = array
# We use the placeholder to assign the variable the intended value.
assign_ops.append(tf.assign(node, placeholder))
saver = tf.train.Saver(nodes_to_save)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
sess.run(assign_ops, feed_dict=feed_dict)
saver.save(sess, checkpoint_path)
assert saver.last_checkpoints[0] == checkpoint_path
def __init__(self, images, labels, fake_data=False, one_hot=False,
dtype=tf.float32):
"""Construct a DataSet.
one_hot arg is used only if fake_data is true. `dtype` can be either
`uint8` to leave the input as `[0, 255]`, or `float32` to rescale into
`[0, 1]`.
"""
dtype = tf.as_dtype(dtype).base_dtype
if dtype not in (tf.uint8, tf.float32):
raise TypeError('Invalid image dtype %r, expected uint8 or float32' %
dtype)
if fake_data:
self._num_examples = 10000
self.one_hot = one_hot
else:
assert images.shape[0] == labels.shape[0], (
'images.shape: %s labels.shape: %s' % (images.shape,
labels.shape))
self._num_examples = images.shape[0]
# Convert shape from [num examples, rows, columns, depth]
# to [num examples, rows*columns] (assuming depth == 1)
assert images.shape[3] == 1
images = images.reshape(images.shape[0],
images.shape[1] * images.shape[2])
if dtype == tf.float32:
# Convert from [0, 255] -> [0.0, 1.0].
images = images.astype(numpy.float32)
images = numpy.multiply(images, 1.0 / 255.0)
self._images = images
self._labels = labels
self._epochs_completed = 0
self._index_in_epoch = 0
def __init__(self,
scale=1.0,
mode='fan_in',
distribution='truncated_normal',
seed=None,
dtype=tf.float32):
if scale <= 0:
raise ValueError('scale must be a positive float: {}'.format(
str(scale)))
if mode not in ['fan_in', 'fan_out', 'fan_avg']:
raise ValueError('invalid mode: {}'.format(str(mode)))
if distribution not in [
'uniform', 'truncated_normal', 'untruncated_normal'
]:
raise ValueError('invalid distribution: {}'.format(
str(distribution)))
if not dtype.is_floating:
raise ValueError('dtype must be a floating point type: {}'.format(
str(dtype)))
self.scale = scale
self.mode = mode
self.distribution = distribution
self.seed = seed
self.dtype = tf.as_dtype(dtype)
def _compute_health_pill(self, x):
x_clean = x[np.where(
np.logical_and(np.logical_not(np.isnan(x)),
np.logical_not(np.isinf(x))))]
if np.size(x_clean):
x_min = np.min(x_clean)
x_max = np.max(x_clean)
x_mean = np.mean(x_clean)
x_var = np.var(x_clean)
else:
x_min = np.inf
x_max = -np.inf
x_mean = np.nan
x_var = np.nan
return np.array([
1.0, # Assume is initialized.
np.size(x),
np.sum(np.isnan(x)),
np.sum(x == -np.inf),
np.sum(np.logical_and(x < 0.0, x != -np.inf)),
np.sum(x == 0.0),
np.sum(np.logical_and(x > 0.0, x != np.inf)),
np.sum(x == np.inf),
x_min,
x_max,
x_mean,
x_var,
float(tf.as_dtype(x.dtype).as_datatype_enum),
float(len(x.shape)),
] + list(x.shape))
def _read_data_types_and_shapes(filenames):
"""Gets dtypes and shapes for all keys in the dataset."""
sequences = _read_numpy_sequences(filenames)
sequence = next(sequences)
sequences.close()
dtypes = {k: tf.as_dtype(v.dtype) for k, v in sequence.items()}
shapes = {k: (None, ) + v.shape[1:] for k, v in sequence.items()}
return dtypes, shapes
def parse_tf_example(tf_example_string, params):
"""Parse the serialized tf.Example and decode it to the image tensor."""
decoded_tensors = tf.parse_single_example(
serialized=tf_example_string,
features={'image/ct_image': tf.FixedLenFeature([], tf.string)})
image = tf.decode_raw(decoded_tensors['image/ct_image'],
tf.as_dtype(tf.float32))
image_size = params.input_image_size + [params.num_channels]
image = tf.reshape(image, image_size)
return image
def cast_clip(clip):
"""
Cast clipping range argument if needed.
"""
if t.dtype in (tf.float32, tf.float64):
if hasattr(clip, 'dtype'):
# Convert to tf dtype in case this is a numpy dtype
clip_dtype = tf.as_dtype(clip.dtype)
if clip_dtype != t.dtype:
return tf.cast(clip, t.dtype)
return clip
def python_type_to_attr_value(
value, #type: Any
attr_name #type: String
):
# type (...) -> tf.AttrValue
"""
Convert a Python object or scalar value to a TensorFlow `tf.AttrValue`
protocol buffer message.
Args:
value: Python object to be converted
Returns:
An AttrValue object that wraps the contents of `value` in the most
appropriate way available.
"""
if isinstance(value, list) or isinstance(value, tuple):
if 0 == len(value):
return tf.AttrValue(list=tf.AttrValue.ListValue())
else:
# Nonempty list
list_value = tf.AttrValue.ListValue()
for elem in value:
# TODO(frreiss): Should we disallow heterogeneous types in lists?
_python_type_to_attr_list_elem(list_value, elem, attr_name)
return tf.AttrValue(list=list_value)
elif isinstance(value, tf.AttrValue):
# TODO(frreiss): Should this case result in an error?
return value
# Scalar types, in the order they appear in the .proto file
elif isinstance(value, string_types):
return tf.AttrValue(s=tf.compat.as_bytes(value))
# Must check for bool before int because bool is a subclass of int in Python
elif isinstance(value, bool):
return tf.AttrValue(b=value)
elif (isinstance(value, int) or isinstance(value, np.int32)
or isinstance(value, np.int64)):
return tf.AttrValue(i=value)
elif isinstance(value, float) or isinstance(value, np.float):
return tf.AttrValue(f=value)
elif isinstance(value, tf.DType):
return tf.AttrValue(type=value.as_datatype_enum)
elif isinstance(value, np.dtype):
return tf.AttrValue(type=tf.as_dtype(value).as_datatype_enum)
elif isinstance(value, tf.TensorShape):
return tf.AttrValue(shape=value.as_proto())
elif isinstance(value, np.ndarray):
return tf.AttrValue(tensor=tf.make_tensor_proto(values=value))
# TODO(frreiss): Populate the "func" and "placeholder" fields of the union
# here
else:
raise ValueError("Don't know how to convert a {} to "
"tf.AttrValue for attribute {}".format(
type(value), attr_name))
def toy_model(features, mesh):
"""A toy model implemented by mesh tensorlfow."""
batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
io_dim = mtf.Dimension('io', FLAGS.io_size)
master_dtype = tf.as_dtype(FLAGS.master_dtype)
slice_dtype = tf.as_dtype(FLAGS.slice_dtype)
activation_dtype = tf.as_dtype(FLAGS.activation_dtype)
x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
x = mtf.cast(x, activation_dtype)
h = x
for lnum in range(1, FLAGS.num_hidden_layers + 2):
if lnum + 1 == FLAGS.num_hidden_layers + 2:
# output layer
dim = io_dim
elif lnum % 2 == 0:
dim = mtf.Dimension('hidden_even', FLAGS.hidden_size)
else:
dim = mtf.Dimension('hidden_odd', FLAGS.hidden_size)
h = mtf.layers.dense(h,
dim,
use_bias=False,
master_dtype=master_dtype,
slice_dtype=slice_dtype,
name='layer_%d' % lnum)
y = h
g = tf.train.get_global_step()
if FLAGS.step_with_nan >= 0:
# Trigger NaN in the forward pass, this is used for testing whether
# MeshTensorFlow can handle occasional NaN value.
y += mtf.import_tf_tensor(
mesh,
tf.divide(
0.0,
tf.cond(tf.equal(g, FLAGS.step_with_nan), lambda: 0.,
lambda: 1.)), mtf.Shape([]))
loss = mtf.reduce_mean(mtf.square(y - x))
return y, loss
def G_mapping(
latents_in, # First input: Latent vectors (Z) [minibatch, latent_size].
labels_in, # Second input: Conditioning labels [minibatch, label_size].
latent_size = 512, # Latent vector (Z) dimensionality.
label_size = 0, # Label dimensionality, 0 if no labels.
dlatent_size = 512, # Disentangled latent (W) dimensionality.
dlatent_broadcast = None, # Output disentangled latent (W) as [minibatch, dlatent_size] or [minibatch, dlatent_broadcast, dlatent_size].
mapping_layers = 8, # Number of mapping layers.
mapping_fmaps = 512, # Number of activations in the mapping layers.
mapping_lrmul = 0.01, # Learning rate multiplier for the mapping layers.
mapping_nonlinearity = 'lrelu', # Activation function: 'relu', 'lrelu'.
use_wscale = True, # Enable equalized learning rate?
normalize_latents = True, # Normalize latent vectors (Z) before feeding them to the mapping layers?
dtype = 'float32', # Data type to use for activations and outputs.
**_kwargs): # Ignore unrecognized keyword args.
act, gain = {'relu': (tf.nn.relu, np.sqrt(2)), 'lrelu': (leaky_relu, np.sqrt(2))}[mapping_nonlinearity]
# Inputs.
latents_in.set_shape([None, latent_size])
labels_in.set_shape([None, label_size])
latents_in = tf.cast(latents_in, dtype)
labels_in = tf.cast(labels_in, dtype)
x = latents_in
# Embed labels and concatenate them with latents.
if label_size:
with tf.variable_scope('LabelConcat'):
w = tf.get_variable('weight', shape=[label_size, latent_size], initializer=tf.initializers.random_normal())
y = tf.matmul(labels_in, tf.cast(w, dtype))
x = tf.concat([x, y], axis=1)
# Normalize latents.
if normalize_latents:
x = pixel_norm(x)
# Mapping layers.
for layer_idx in range(mapping_layers):
with tf.variable_scope('Dense%d' % layer_idx):
fmaps = dlatent_size if layer_idx == mapping_layers - 1 else mapping_fmaps
x = dense(x, fmaps=fmaps, gain=gain, use_wscale=use_wscale, lrmul=mapping_lrmul)
x = apply_bias(x, lrmul=mapping_lrmul)
x = act(x)
# Broadcast.
if dlatent_broadcast is not None:
with tf.variable_scope('Broadcast'):
x = tf.tile(x[:, np.newaxis], [1, dlatent_broadcast, 1])
# Output.
assert x.dtype == tf.as_dtype(dtype)
return tf.identity(x, name='dlatents_out')
def _decode_liver_example(serialized_example):
"""Parses a single tf.Example into image and label tensors."""
features = {}
features['image/ct_image'] = tf.FixedLenFeature([], tf.string)
features['image/label'] = tf.FixedLenFeature([], tf.string)
parsed = tf.parse_single_example(serialized_example,
features=features)
# Here, assumes the `image` is normalized to [0, 1] of type float32 and
# the `label` is a binary matrix, whose last dimension is one_hot encoded
# labels.
# The dtype of `label` can be either float32 or int64.
image = tf.decode_raw(parsed['image/ct_image'],
tf.as_dtype(tf.float32))
label = tf.decode_raw(parsed['image/label'],
tf.as_dtype(params.label_dtype))
image_size = params.input_image_size + [params.num_channels]
image = tf.reshape(image, image_size)
label_size = params.input_image_size + [params.num_classes]
label = tf.reshape(label, label_size)
if self._is_training and params.use_index_label_in_train:
# Use class index for labels and remove the channel dim (#channels=1).
channel_dim = -1
label = tf.argmax(label,
axis=channel_dim,
output_type=tf.int32)
if params.use_bfloat16:
image = tf.cast(image, dtype=tf.bfloat16)
if label.dtype == tf.float32:
label = tf.cast(label, dtype=tf.bfloat16)
# TPU doesn't support tf.int64 well, use tf.int32 directly.
if label.dtype == tf.int64:
label = tf.cast(label, dtype=tf.int32)
return image, label
def create_input_variable(self, input):
for i in range(len(input)):
print(input[i].shape)
placeholder = tf.compat.v1.placeholder(tf.as_dtype(input[i].dtype),
shape=input[i].shape)
var = tf.Variable(
placeholder,
trainable=False,
collections=[tf.compat.v1.GraphKeys.LOCAL_VARIABLES])
self.variable_initialization[placeholder] = input[i]
input[i] = var
for i in input:
print(i.shape)
return input
def load_dataset_from_directory(directory,
length,
batch,
cache_update_every=1000,
buffer_size=10):
loader = functools.partial(numpy_loader, directory, cache_update_every)
dtypes, shapes = read_spec(loader)
dtypes = {key: tf.as_dtype(value) for key, value in dtypes.items()}
shapes = {key: (None,) + shape[1:] for key, shape in shapes.items()}
chunking = functools.partial(chunk_sequence, length=length)
dataset = tfdd.from_generator(loader, dtypes, shapes)
dataset = dataset.flat_map(chunking)
dataset = dataset.batch(batch, drop_remainder=True)
dataset = dataset.prefetch(buffer_size)
return dataset
def getOrCreateSparseLinear(self,
x_shape,
x_dtype,
sparsity,
dense_length,
block_size,
use_bias,
override_partials_type=None):
x_dtype = tf.as_dtype(x_dtype)
# Each layer should have a unique scope name
scope_name = tf.get_default_graph().get_name_scope()
logger.info(f"Sparse layer with scope name: {scope_name}")
# Construct the layer if it does not exist
if scope_name not in self.sparse_layers:
layer_matmul_options = self.sparse_matmul_options
if override_partials_type:
layer_matmul_options['partialsType'] = override_partials_type
limit = np.sqrt(6 / ((x_shape[-1] + dense_length) *
(1 - sparsity)))
uniform_gen = partial(self.random.uniform, -limit, limit)
indices_random_gen = np.random.default_rng(seed=self.random_seed)
sparse_layer = layers.SparseFcLayer.from_random_generator(
dense_length,
x_shape,
1 - sparsity,
block_size=block_size,
values_initialiser_gen=uniform_gen,
indices_initialiser_gen=indices_random_gen,
name="sparse_layer",
dtype=x_dtype,
matmul_options=layer_matmul_options,
use_bias=use_bias,
relu=False,
disable_updating=self.disable_updating,
pooling_type=self.pooling_type)
# Create placeholders on the host, outside XLA
with tf.init_scope(): # escapes XLA
with tf.device("cpu"):
sparse_layer.create_placeholders()
self.sparse_layers[scope_name] = sparse_layer
else:
# Re-use a previously defined layer
sparse_layer = self.sparse_layers[scope_name]
return sparse_layer
def irdft_matrix(shape, dtype=tf.float32):
"""Matrix for implementing kernel reparameterization with `tf.matmul`.
This can be used to represent a kernel with the provided shape in the RDFT
domain.
Example code for kernel creation, assuming 2D kernels:
```
def create_kernel(init):
shape = init.shape.as_list()
matrix = irdft_matrix(shape[:2])
init = tf.reshape(init, (shape[0] * shape[1], shape[2] * shape[3]))
init = tf.matmul(tf.transpose(matrix), init)
kernel = tf.Variable(init)
kernel = tf.matmul(matrix, kernel)
kernel = tf.reshape(kernel, shape)
return kernel
```
Args:
shape: Iterable of integers. Shape of kernel to apply this matrix to.
dtype: `dtype` of returned matrix.
Returns:
`Tensor` of shape `(prod(shape), prod(shape))` and dtype `dtype`.
"""
shape = tuple(int(s) for s in shape)
dtype = tf.as_dtype(dtype)
size = np.prod(shape)
rank = len(shape)
matrix = np.identity(size, dtype=np.float64).reshape((size, ) + shape)
for axis in range(rank):
matrix = fftpack.rfft(matrix, axis=axis + 1)
slices = (rank + 1) * [slice(None)]
if shape[axis] % 2 == 1:
slices[axis + 1] = slice(1, None)
else:
slices[axis + 1] = slice(1, -1)
matrix[tuple(slices)] *= np.sqrt(2)
matrix /= np.sqrt(size)
matrix = np.reshape(matrix, (size, size))
return tf.constant(matrix,
dtype=dtype,
name="irdft_" + "x".join([str(s) for s in shape]))
def Test():
x = tf.placeholder(dtype=tf.float32, shape=[None])
batch_size = tf.shape(x)[0]
r = tf.convert_to_tensor([batch_size, 1])
tensor_info_x = meta_graph_pb2.TensorInfo(name=x.name,
dtype=tf.as_dtype(
x.dtype).as_datatype_enum)
tensor_info_r = tf.compat.v1.saved_model.utils.build_tensor_info(r)
return {
'key':
(tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
inputs={'x': tensor_info_x},
outputs={'r': tensor_info_r},
method_name='some_function'))
}, None, None
def __init__(self,
symbols_to_logits_fn,
vocab_size,
batch_size,
beam_size,
alpha,
max_decode_length,
eos_id,
padded_decode,
dtype=tf.float32):
"""Initialize sequence beam search.
Args:
symbols_to_logits_fn: A function to provide logits, which is the
interface to the Transformer model. The passed in arguments are:
ids -> A tensor with shape [batch_size * beam_size, index].
index -> A scalar.
cache -> A nest dictionary of tensors [batch_size * beam_size, ...].
The function must return a tuple of logits and the updated cache:
logits -> A tensor with shape [batch * beam_size, vocab_size].
updated cache -> A nested dictionary with the same structure as the
input cache.
vocab_size: An integer, the size of the vocabulary, used for topk
computation.
batch_size: An integer, the decode batch size.
beam_size: An integer, number of beams for beam search.
alpha: A float, defining the strength of length normalization.
max_decode_length: An integer, the maximum number of steps to decode
a sequence.
eos_id: An integer. ID of end of sentence token.
padded_decode: A bool, indicating if max_sequence_length padding is used
for beam search.
dtype: A tensorflow data type used for score computation. The default is
tf.float32.
"""
self.symbols_to_logits_fn = symbols_to_logits_fn
self.vocab_size = vocab_size
self.batch_size = batch_size
self.beam_size = beam_size
self.alpha = alpha
self.max_decode_length = max_decode_length
self.eos_id = eos_id
self.padded_decode = padded_decode
self.dtype = tf.as_dtype(dtype)
def _set_params_initializer(self, hparams, mode, features):
"""Set various params for self and initialize."""
self.mode = mode
self.src_vocab_size = hparams.src_vocab_size
self.tgt_vocab_size = hparams.tgt_vocab_size
self.features = features
self.dtype = tf.as_dtype(hparams.activation_dtype)
self.single_cell_fn = None
# Set num units
self.num_units = hparams.num_units
self.eos_id = hparams.tgt_eos_id
self.label_smoothing = hparams.label_smoothing
# Set num layers
self.num_encoder_layers = hparams.num_encoder_layers
self.num_decoder_layers = hparams.num_decoder_layers
assert self.num_encoder_layers
assert self.num_decoder_layers
# Batch size
self.batch_size = tf.size(self.features["source_sequence_length"])
# Global step
# Use get_global_step instead of user-defied global steps. Otherwise the
# num_train_steps in TPUEstimator.train has no effect (will train forever).
# TPUestimator only check if tf.train.get_global_step() < num_train_steps
self.global_step = tf.train.get_or_create_global_step()
# Initializer
self.random_seed = hparams.random_seed
initializer = model_helper.get_initializer(hparams.init_op,
self.random_seed,
hparams.init_weight)
tf.get_variable_scope().set_initializer(initializer)
# Embeddings
self.encoder_emb_lookup_fn = (self._emb_lookup if self.mode
== contrib_learn.ModeKeys.TRAIN else
tf.nn.embedding_lookup)
def attention(theta, new_lstm_state):
"""Helper function to add attention."""
lstm_output = new_lstm_state["h"]
query = tf.expand_dims(tf.matmul(lstm_output, theta["query_kernel"]), 0)
normed_v = theta["atten_g"] * theta["atten_v"] * tf.rsqrt(
tf.reduce_sum(tf.square(theta["atten_v"])))
score = tf.reduce_sum(
normed_v * tf.tanh(theta["keys"] + query + theta["atten_b"]), [2])
score = tf.transpose(score)
score = tf.where(
theta["seq_mask"] > 0.5, score,
tf.ones_like(score) * tf.as_dtype(score.dtype).as_numpy_dtype(-np.inf))
alignments = tf.nn.softmax(score)
score = tf.transpose(alignments)
atten = tf.reduce_sum(tf.expand_dims(score, 2) * theta["values"], 0)
new_states = {"attention": atten, "alignments": alignments}
for k in new_lstm_state:
new_states[k] = new_lstm_state[k]
return new_states
def _get_stacked_2d_slices(image_3d, label_3d):
"""Return 2d slices of the 3d scan."""
image_stack = []
label_stack = []
for begin_idx in range(0, FLAGS.ct_resolution - FLAGS.image_c + 1):
slice_begin = [0, 0, begin_idx]
slice_size = [
FLAGS.ct_resolution, FLAGS.ct_resolution, FLAGS.image_c
]
image = tf.slice(image_3d, slice_begin, slice_size)
slice_begin = [0, 0, begin_idx + FLAGS.image_c // 2]
slice_size = [FLAGS.ct_resolution, FLAGS.ct_resolution, 1]
label = tf.slice(label_3d, slice_begin, slice_size)
spatial_dims_w_blocks = [
FLAGS.image_nx_block,
FLAGS.ct_resolution // FLAGS.image_nx_block,
FLAGS.image_ny_block,
FLAGS.ct_resolution // FLAGS.image_ny_block
]
image = tf.reshape(image,
spatial_dims_w_blocks + [FLAGS.image_c])
label = tf.reshape(label, spatial_dims_w_blocks)
label = tf.cast(label, tf.int32)
label = tf.one_hot(label, FLAGS.label_c)
data_dtype = tf.as_dtype(FLAGS.mtf_dtype)
image = tf.cast(image, data_dtype)
label = tf.cast(label, data_dtype)
image_stack.append(image)
label_stack.append(label)
return tf.stack(image_stack), tf.stack(label_stack)
def build(self, input_shape):
"""Implements build() for the layer."""
dtype = tf.as_dtype(self.dtype or tf.keras.backend.floatx())
if not (dtype.is_floating or dtype.is_complex):
raise TypeError(
"Unable to build `Dense3D` layer with non-floating "
"point (and non-complex) dtype %s" % (dtype, ))
input_shape = tf.TensorShape(input_shape)
if tf.compat.dimension_value(input_shape[-1]) is None:
raise ValueError("The last dimension of the inputs to `Dense3D` "
"should be defined. Found `None`.")
self.last_dim = tf.compat.dimension_value(input_shape[-1])
self.input_spec = tf.keras.layers.InputSpec(min_ndim=3,
axes={-1: self.last_dim})
# Determines variable shapes.
if self.backward_compatible:
kernel_shape = self.compatible_kernel_shape
bias_shape = self.compatible_bias_shape
else:
kernel_shape = self.kernel_shape
bias_shape = self.bias_shape
self.kernel = self.add_weight("kernel",
shape=kernel_shape,
initializer=self.kernel_initializer,
dtype=self.dtype,
trainable=True)
if self.use_bias:
self.bias = self.add_weight("bias",
shape=bias_shape,
initializer=self.bias_initializer,
dtype=self.dtype,
trainable=True)
else:
self.bias = None
super(Dense3D, self).build(input_shape)
def _get_tf_dtype(space):
if isinstance(space, gym.spaces.Discrete):
return tf.int32
if isinstance(space, gym.spaces.Box):
return tf.as_dtype(space.dtype)
raise NotImplementedError()
def activation_dtype(self): return tf.as_dtype(self._hparams.activation_dtype)
def slice_dtype(self): return tf.as_dtype(self._hparams.slice_dtype)
