def make_status_message(model): """Makes a string `Tensor` of training status.""" return tf.strings.join([ 'Starting train step: current_image_id: ', tf.as_string(model.current_image_id), ', progress: ', tf.as_string(model.progress), ', num_blocks: {}'.format( model.num_blocks), ', batch_size: {}'.format(model.batch_size) ], name='status_message')
def op(name, images, max_outputs=3, display_name=None, description=None, collections=None): """Create a legacy image summary op for use in a TensorFlow graph. Arguments: name: A unique name for the generated summary node. images: A `Tensor` representing pixel data with shape `[k, h, w, c]`, where `k` is the number of images, `h` and `w` are the height and width of the images, and `c` is the number of channels, which should be 1, 3, or 4. Any of the dimensions may be statically unknown (i.e., `None`). max_outputs: Optional `int` or rank-0 integer `Tensor`. At most this many images will be emitted at each step. When more than `max_outputs` many images are provided, the first `max_outputs` many images will be used and the rest silently discarded. display_name: Optional name for this summary in TensorBoard, as a constant `str`. Defaults to `name`. description: Optional long-form description for this summary, as a constant `str`. Markdown is supported. Defaults to empty. collections: Optional list of graph collections keys. The new summary op is added to these collections. Defaults to `[Graph Keys.SUMMARIES]`. Returns: A TensorFlow summary op. """ # TODO(nickfelt): remove on-demand imports once dep situation is fixed. import tensorflow.compat.v1 as tf if display_name is None: display_name = name summary_metadata = metadata.create_summary_metadata( display_name=display_name, description=description) with tf.name_scope(name), \ tf.control_dependencies([tf.assert_rank(images, 4), tf.assert_type(images, tf.uint8), tf.assert_non_negative(max_outputs)]): limited_images = images[:max_outputs] encoded_images = tf.map_fn(tf.image.encode_png, limited_images, dtype=tf.string, name='encode_each_image') image_shape = tf.shape(images) dimensions = tf.stack([ tf.as_string(image_shape[2], name='width'), tf.as_string(image_shape[1], name='height') ], name='dimensions') tensor = tf.concat([dimensions, encoded_images], axis=0) return tf.summary.tensor_summary(name='image_summary', tensor=tensor, collections=collections, summary_metadata=summary_metadata)
def _extend_with_dummy(extend_with, to_extend, dummy_value='n/a'): """Extends one SparseTensor with dummy_values at positions of other.""" dense_shape = tf.to_int64( tf.concat([[tf.shape(extend_with)[0]], [tf.maximum(tf.shape(extend_with)[1], tf.shape(to_extend)[1])], [tf.maximum(tf.shape(extend_with)[2], tf.shape(to_extend)[2])]], axis=0)) additional_indices = tf.sets.set_difference( tf.SparseTensor( indices=extend_with.indices, values=tf.zeros_like(extend_with.values, dtype=tf.int32), dense_shape=dense_shape), tf.SparseTensor( indices=to_extend.indices, values=tf.zeros([tf.shape(to_extend.indices)[0]], dtype=tf.int32), dense_shape=dense_shape)).indices # Supply defaults for all other indices. default = tf.tile( tf.constant([dummy_value]), multiples=[tf.shape(additional_indices)[0]]) string_value = ( tf.as_string(to_extend.values) if to_extend.values.dtype != tf.string else to_extend.values) return tf.sparse_reorder( tf.SparseTensor( indices=tf.concat([to_extend.indices, additional_indices], axis=0), values=tf.concat([string_value, default], axis=0), dense_shape=dense_shape))
def _replace_empty_string_with_random_number(string_tensor): """Returns string unchanged if non-empty, and random string tensor otherwise. The random string is an integer 0 and 2**63 - 1, casted as string. Args: string_tensor: A tf.tensor of dtype string. Returns: out_string: A tf.tensor of dtype string. If string_tensor contains the empty string, out_string will contain a random integer casted to a string. Otherwise string_tensor is returned unchanged. """ empty_string = tf.constant('', dtype=tf.string, name='EmptyString') random_source_id = tf.as_string( tf.random_uniform(shape=[], maxval=2**63 - 1, dtype=tf.int64)) out_string = tf.cond( tf.equal(string_tensor, empty_string), true_fn=lambda: random_source_id, false_fn=lambda: string_tensor) return out_string
def sampled_softmax_loss(src_emb, pos_ids, neg_num, output_emb_table, output_emb_bias, node_size, s2h=True): """Sampled softmax loss. Args: src_emb: positive src embedding with shape [batch_size, dim] pos_ids: positive ids. output_emb_table: output_emb_bias: node_size: total node size. s2h: set True if need string to hash. """ if s2h: pos_ids = tf.as_string(pos_ids) pos_ids = tf.string_to_hash_bucket_fast( pos_ids, node_size, name='softmax_loss_to_hash_bucket_oper') loss = tf.nn.sampled_softmax_loss( weights=output_emb_table, biases=output_emb_bias, labels=tf.reshape(pos_ids, [-1, 1]), inputs=src_emb, num_sampled=neg_num, num_classes=node_size, partition_strategy='mod', remove_accidental_hits=True) return [tf.reduce_mean(loss), None, None]
def _dedup_tensor(sp_tensor: tf.SparseTensor) -> tf.SparseTensor: """Dedup values of a SparseTensor along each row. Args: sp_tensor: A 2D SparseTensor to be deduped. Returns: A deduped SparseTensor of shape [batch_size, max_len], where max_len is the maximum number of unique values for a row in the Tensor. """ string_batch_index = tf.as_string(sp_tensor.indices[:, 0]) # tf.unique only works on 1D tensors. To avoid deduping across examples, # prepend each feature value with the example index. This requires casting # to and from strings for non-string features. string_values = sp_tensor.values original_dtype = sp_tensor.values.dtype if original_dtype != tf.string: string_values = tf.as_string(sp_tensor.values) index_and_value = tf.strings.join([string_batch_index, string_values], separator='|') unique_index_and_value, _ = tf.unique(index_and_value) # split is a shape [tf.size(values), 2] tensor. The first column contains # indices and the second column contains the feature value (we assume no # feature contains | so we get exactly 2 values from the string split). split = tf.string_split(unique_index_and_value, delimiter='|') split = tf.reshape(split.values, [-1, 2]) string_indices = split[:, 0] values = split[:, 1] indices = tf.reshape( tf.string_to_number(string_indices, out_type=tf.int32), [-1]) if original_dtype != tf.string: values = tf.string_to_number(values, out_type=original_dtype) values = tf.reshape(values, [-1]) # Convert example indices into SparseTensor indices, e.g. # [0, 0, 0, 1, 3, 3] -> [[0,0], [0,1], [0,2], [1,0], [3,0], [3,1]] batch_size = tf.to_int32(sp_tensor.dense_shape[0]) new_indices, max_len = _example_index_to_sparse_index(indices, batch_size) return tf.SparseTensor( indices=tf.to_int64(new_indices), values=values, dense_shape=[tf.to_int64(batch_size), max_len])
def train(hparams): """Trains a CycleGAN. Args: hparams: An HParams instance containing the hyperparameters for training. """ if not tf.io.gfile.exists(hparams.train_log_dir): tf.io.gfile.makedirs(hparams.train_log_dir) with tf.device(tf.train.replica_device_setter(hparams.ps_replicas)): with tf.name_scope('inputs'), tf.device('/cpu:0'): images_x, images_y = _get_data(hparams.image_set_x_file_pattern, hparams.image_set_y_file_pattern, hparams.batch_size, hparams.patch_size) # Define CycleGAN model. cyclegan_model = _define_model(images_x, images_y) # Define CycleGAN loss. cyclegan_loss = tfgan.cyclegan_loss( cyclegan_model, cycle_consistency_loss_weight=hparams. cycle_consistency_loss_weight, tensor_pool_fn=tfgan.features.tensor_pool) # Define CycleGAN train ops. train_ops = _define_train_ops(cyclegan_model, cyclegan_loss, hparams) # Training train_steps = tfgan.GANTrainSteps(1, 1) status_message = tf.strings.join([ 'Starting train step: ', tf.as_string(tf.train.get_or_create_global_step()) ], name='status_message') if not hparams.max_number_of_steps: return tfgan.gan_train( train_ops, hparams.train_log_dir, get_hooks_fn=tfgan.get_sequential_train_hooks(train_steps), hooks=[ tf.estimator.StopAtStepHook( num_steps=hparams.max_number_of_steps), tf.estimator.LoggingTensorHook( {'status_message': status_message}, every_n_iter=10) ], master=hparams.master, is_chief=hparams.task == 0)
def encode(self, ego_tensor): if not self._use_edge: ids = ego_tensor.src.ids else: # TODO: replace this hack with edge ids. ids = tf.cast(ego_tensor.src.continuous_attrs, dtype=tf.int64) if self._str2hash: index = tf.as_string(ids) ids = tf.string_to_hash_bucket_fast(index, self._num, name=self._name + 'str_to_hash_bucket_op') emb = tf.nn.embedding_lookup(self._emb_table, ids, name=self._name + 'ids_embedding_lookup_op') emb = tf.reshape(emb, [-1, self._dim]) return emb
# A quick fix to run TF 1.X code in TF 2.X, we may want to properly migrate the Python script to TF 2.X API. import tensorflow.compat.v1 as tf tf.disable_v2_behavior() import os a = tf.placeholder(tf.float32, shape=None, name="a") b = tf.placeholder(tf.float32, shape=None, name="b") e = tf.placeholder(tf.string, shape=None, name="e") ee = tf.strings.to_number(e, out_type=tf.float32, name="ee") v = tf.Variable(dtype=tf.float32, initial_value=tf.constant(1.0), name="v") c = tf.add(a, b, name="c") d = tf.add(c, v) eee = tf.add(ee, v) e4 = tf.as_string(eee) global_step = tf.train.get_or_create_global_step() global_step_inc = tf.assign_add(global_step, 1) hooks = [tf.train.StopAtStepHook(last_step=2)] with tf.Session() as mon_sess: mon_sess.run(tf.initialize_all_variables()) for i in range(2): print( mon_sess.run([d, e4, global_step_inc], feed_dict={ a: [1.0, 2.0, 3.0], b: [1.0, 2.0, 3.0], e: ["1.0", "2.0", "3.0"] })) signatures = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: { 'inputs': {
def train(hparams): """Trains a StarGAN. Args: hparams: An HParams instance containing the hyperparameters for training. """ # Create the log_dir if not exist. if not tf.io.gfile.exists(hparams.train_log_dir): tf.io.gfile.makedirs(hparams.train_log_dir) # Shard the model to different parameter servers. with tf.device(tf.train.replica_device_setter(hparams.ps_replicas)): # Create the input dataset. with tf.name_scope('inputs'), tf.device('/cpu:0'): images, labels = data_provider.provide_data( 'train', hparams.batch_size, hparams.patch_size) # Define the model. with tf.name_scope('model'): model = _define_model(images, labels) # Add image summary. tfgan.eval.add_stargan_image_summaries(model, num_images=3 * hparams.batch_size, display_diffs=True) # Define the model loss. loss = tfgan.stargan_loss(model) # Define the train ops. with tf.name_scope('train_ops'): train_ops = _define_train_ops(model, loss, hparams.generator_lr, hparams.discriminator_lr, hparams.adam_beta1, hparams.adam_beta2, hparams.max_number_of_steps) # Define the train steps. train_steps = _define_train_step(hparams.gen_disc_step_ratio) # Define a status message. status_message = tf.strings.join([ 'Starting train step: ', tf.as_string(tf.train.get_or_create_global_step()) ], name='status_message') # Train the model. tfgan.gan_train( train_ops, hparams.train_log_dir, get_hooks_fn=tfgan.get_sequential_train_hooks(train_steps), hooks=[ tf.estimator.StopAtStepHook( num_steps=hparams.max_number_of_steps), tf.estimator.LoggingTensorHook([status_message], every_n_iter=10) ], master=hparams.tf_master, is_chief=hparams.task == 0)
def train(hparams): """Trains an MNIST GAN. Args: hparams: An HParams instance containing the hyperparameters for training. """ if not tf.io.gfile.exists(hparams.train_log_dir): tf.io.gfile.makedirs(hparams.train_log_dir) # Force all input processing onto CPU in order to reserve the GPU for # the forward inference and back-propagation. with tf.name_scope('inputs'), tf.device('/cpu:0'): images, one_hot_labels = data_provider.provide_data( 'train', hparams.batch_size, num_parallel_calls=4) # Define the GANModel tuple. Optionally, condition the GAN on the label or # use an InfoGAN to learn a latent representation. if hparams.gan_type == 'unconditional': gan_model = tfgan.gan_model( generator_fn=networks.unconditional_generator, discriminator_fn=networks.unconditional_discriminator, real_data=images, generator_inputs=tf.random.normal( [hparams.batch_size, hparams.noise_dims])) elif hparams.gan_type == 'conditional': noise = tf.random.normal([hparams.batch_size, hparams.noise_dims]) gan_model = tfgan.gan_model( generator_fn=networks.conditional_generator, discriminator_fn=networks.conditional_discriminator, real_data=images, generator_inputs=(noise, one_hot_labels)) elif hparams.gan_type == 'infogan': cat_dim, cont_dim = 10, 2 generator_fn = functools.partial(networks.infogan_generator, categorical_dim=cat_dim) discriminator_fn = functools.partial(networks.infogan_discriminator, categorical_dim=cat_dim, continuous_dim=cont_dim) unstructured_inputs, structured_inputs = util.get_infogan_noise( hparams.batch_size, cat_dim, cont_dim, hparams.noise_dims) gan_model = tfgan.infogan_model( generator_fn=generator_fn, discriminator_fn=discriminator_fn, real_data=images, unstructured_generator_inputs=unstructured_inputs, structured_generator_inputs=structured_inputs) tfgan.eval.add_gan_model_image_summaries(gan_model, hparams.grid_size) # Get the GANLoss tuple. You can pass a custom function, use one of the # already-implemented losses from the losses library, or use the defaults. with tf.name_scope('loss'): if hparams.gan_type == 'infogan': gan_loss = tfgan.gan_loss( gan_model, generator_loss_fn=tfgan.losses.modified_generator_loss, discriminator_loss_fn=tfgan.losses.modified_discriminator_loss, mutual_information_penalty_weight=1.0, add_summaries=True) else: gan_loss = tfgan.gan_loss(gan_model, add_summaries=True) tfgan.eval.add_regularization_loss_summaries(gan_model) # Get the GANTrain ops using custom optimizers. with tf.name_scope('train'): gen_lr, dis_lr = _learning_rate(hparams.gan_type) train_ops = tfgan.gan_train_ops( gan_model, gan_loss, generator_optimizer=tf.train.AdamOptimizer(gen_lr, 0.5), discriminator_optimizer=tf.train.AdamOptimizer(dis_lr, 0.5), summarize_gradients=True, aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N) # Run the alternating training loop. Skip it if no steps should be taken # (used for graph construction tests). status_message = tf.strings.join([ 'Starting train step: ', tf.as_string(tf.train.get_or_create_global_step()) ], name='status_message') if hparams.max_number_of_steps == 0: return tfgan.gan_train( train_ops, hooks=[ tf.estimator.StopAtStepHook(num_steps=hparams.max_number_of_steps), tf.estimator.LoggingTensorHook([status_message], every_n_iter=10) ], logdir=hparams.train_log_dir, get_hooks_fn=tfgan.get_joint_train_hooks(), save_checkpoint_secs=60)
def _process(examples): """Supplies input to our model. This function supplies input to our model after parsing. Args: examples: The dictionary from key to (Sparse)Tensors with context and sequence features. Returns: A tuple consisting of 1) a dictionary of tensors whose keys are the feature names, and 2) a tensor of target labels if the mode is not INFER (and None, otherwise). """ # Combine into a single dictionary. feature_map = {} # Add age if requested. if include_age: age_in_seconds = ( examples[CONTEXT_KEY_PREFIX + 'timestamp'] - examples.pop(CONTEXT_KEY_PREFIX + 'Patient.birthDate')) age_in_years = tf.to_float(age_in_seconds) / (60 * 60 * 24 * 365.0) feature_map[CONTEXT_KEY_PREFIX + AGE_KEY] = age_in_years sequence_length = examples.pop(CONTEXT_KEY_PREFIX + 'sequenceLength') # Cross the requested features. for cross in time_crossed_features: # The features may be missing at different rates - we take the union # of the indices supplying defaults. extended_features = dict() dense_shape = tf.concat( [[tf.to_int64(tf.shape(sequence_length)[0])], [tf.reduce_max(sequence_length)], tf.constant([1], dtype=tf.int64)], axis=0) for i, feature in enumerate(cross): sp_tensor = examples[SEQUENCE_KEY_PREFIX + feature] additional_indices = [] covered_indices = sp_tensor.indices for j, other_feature in enumerate(cross): if i != j: additional_indices.append( tf.sets.set_difference( tf.sparse_reorder( tf.SparseTensor( indices=examples[ SEQUENCE_KEY_PREFIX + other_feature].indices, values=tf.zeros([ tf.shape(examples[ SEQUENCE_KEY_PREFIX + other_feature].indices)[0] ], dtype=tf.int32), dense_shape=dense_shape)), tf.sparse_reorder( tf.SparseTensor( indices=covered_indices, values=tf.zeros( [tf.shape(covered_indices)[0]], dtype=tf.int32), dense_shape=dense_shape))).indices) covered_indices = tf.concat([sp_tensor.indices] + additional_indices, axis=0) additional_indices = tf.concat(additional_indices, axis=0) # Supply defaults for all other indices. default = tf.tile(tf.constant(['n/a']), multiples=[tf.shape(additional_indices)[0]]) string_value = sp_tensor.values if string_value.dtype != tf.string: string_value = tf.as_string(string_value) extended_features[feature] = tf.sparse_reorder( tf.SparseTensor(indices=tf.concat( [sp_tensor.indices, additional_indices], axis=0), values=tf.concat([string_value, default], axis=0), dense_shape=dense_shape)) new_values = tf.strings.join( [extended_features[f].values for f in cross], separator='-') crossed_sp_tensor = tf.sparse_reorder( tf.SparseTensor( indices=extended_features[cross[0]].indices, values=new_values, dense_shape=extended_features[cross[0]].dense_shape)) examples[SEQUENCE_KEY_PREFIX + '_'.join(cross)] = crossed_sp_tensor # Remove unwanted features that are used in the cross but should not be # considered outside the cross. for cross in time_crossed_features: for feature in cross: if (feature not in sequence_features and SEQUENCE_KEY_PREFIX + feature in examples): del examples[SEQUENCE_KEY_PREFIX + feature] # Flatten sparse tensor to compute event age. This dense tensor also # contains padded values. These will not be used when gathering elements # from the dense tensor since each sparse feature won't have a value # defined for the padding. padded_event_age = ( # Broadcast current time along sequence dimension. tf.expand_dims(examples.pop(CONTEXT_KEY_PREFIX + 'timestamp'), 1) # Subtract time of events. - examples.pop(SEQUENCE_KEY_PREFIX + 'eventId')) for i in range(len(time_windows) - 1): max_age = time_windows[i] min_age = time_windows[i + 1] padded_in_time_window = tf.logical_and(padded_event_age <= max_age, padded_event_age > min_age) for k, v in examples.iteritems(): if k.startswith(CONTEXT_KEY_PREFIX): continue # For each sparse feature entry, look up whether it is in the time # window or not. in_time_window = tf.gather_nd(padded_in_time_window, v.indices[:, 0:2]) v = tf.sparse_retain(v, in_time_window) sp_tensor = tf.sparse_reshape(v, [v.dense_shape[0], -1]) if dedup: sp_tensor = _dedup_tensor(sp_tensor) feature_map[k + '-til-%d' % min_age] = sp_tensor for k, v in examples.iteritems(): if k.startswith(CONTEXT_KEY_PREFIX): feature_map[k] = v return feature_map
def train(hparams): """Trains a CIFAR10 GAN. Args: hparams: An HParams instance containing the hyperparameters for training. """ if not tf.io.gfile.exists(hparams.train_log_dir): tf.io.gfile.makedirs(hparams.train_log_dir) with tf.device(tf.train.replica_device_setter(hparams.ps_replicas)): # Force all input processing onto CPU in order to reserve the GPU for # the forward inference and back-propagation. with tf.name_scope('inputs'): with tf.device('/cpu:0'): images, _ = data_provider.provide_data('train', hparams.batch_size, num_parallel_calls=4) # Define the GANModel tuple. generator_fn = networks.generator discriminator_fn = networks.discriminator generator_inputs = tf.random.normal([hparams.batch_size, 64]) gan_model = tfgan.gan_model(generator_fn, discriminator_fn, real_data=images, generator_inputs=generator_inputs) tfgan.eval.add_gan_model_image_summaries(gan_model) # Get the GANLoss tuple. Use the selected GAN loss functions. with tf.name_scope('loss'): gan_loss = tfgan.gan_loss(gan_model, gradient_penalty_weight=1.0, add_summaries=True) # Get the GANTrain ops using the custom optimizers and optional # discriminator weight clipping. with tf.name_scope('train'): gen_opt, dis_opt = _get_optimizers(hparams) train_ops = tfgan.gan_train_ops(gan_model, gan_loss, generator_optimizer=gen_opt, discriminator_optimizer=dis_opt, summarize_gradients=True) # Run the alternating training loop. Skip it if no steps should be taken # (used for graph construction tests). status_message = tf.strings.join([ 'Starting train step: ', tf.as_string(tf.train.get_or_create_global_step()) ], name='status_message') if hparams.max_number_of_steps == 0: return tfgan.gan_train(train_ops, hooks=([ tf.estimator.StopAtStepHook( num_steps=hparams.max_number_of_steps), tf.estimator.LoggingTensorHook([status_message], every_n_iter=10) ]), logdir=hparams.train_log_dir, master=hparams.master, is_chief=hparams.task == 0)