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)