def preprocess(videos, target_resolution): """Runs some preprocessing on the videos for I3D model. Args: videos: <T>[batch_size, num_frames, height, width, depth] The videos to be preprocessed. We don't care about the specific dtype of the videos, it can be anything that tf.image.resize_bilinear accepts. Values are expected to be in the range 0-255. target_resolution: (width, height): target video resolution Returns: videos: <float32>[batch_size, num_frames, height, width, depth] """ videos_shape = videos.shape.as_list() all_frames = tf.reshape(videos, [-1] + videos_shape[-3:]) resized_videos = tf.image.resize_bilinear(all_frames, size=target_resolution) target_shape = [videos_shape[0], -1] + list(target_resolution) + [3] output_videos = tf.reshape(resized_videos, target_shape) ops = [ tf.print("output_videos Max:", tf.reduce_max(output_videos)), tf.print("output_videos Min:", tf.reduce_min(output_videos)) ] #scaled_videos = 2. * tf.cast(output_videos, tf.float32) / 255. - 1 scaled_videos = 2. * output_videos / 255. - 1 return scaled_videos
def reverb_dataset_fn(i): tf.print('Creating dataset for replica; index:', i) return reverb_dataset.ReplayDataset( self._client.server_address, table=tf.constant('dist'), dtypes=(tf.float32,), shapes=(tf.TensorShape([3, 3]),), max_in_flight_samples_per_worker=100).take(2)
def check_probabilities(_, v): probability = v.info.probability self.assertLen(probability.values, 4) # Don't use any math ops since tensor values seem to contain # unaligned tensors on some systems; but tf.print doesn't check alignment. # # This seems to be caused by a compatibility issue where DistStrat isn't # well tested when eager mode is disabled. So instead of treating this # as a true TF bug, we just work around it. We can remove this hack and # convert it to e.g. tf.assert_greater type check if/when we enable eager # execution for these tests. tf.print('Probability values:', probability.values)
def update(self, samples, contexts, dev_samples, dev_contexts): if self._counter % 20 == 0: # To prevent memory leaks in tf eager tf.set_random_seed(self._seed) actions, rews, weights, kwargs = self.create_batch( samples, contexts=contexts) dev_actions, dev_rews, dev_weights, dev_kwargs = self.create_batch( dev_samples, contexts=dev_contexts) trajs = (s.traj for s in samples) with tf.GradientTape( watch_accessed_variables=False, persistent=True) as tape0: tape0.watch(self._score_vars) scores = self.compute_scores(trajs, return_tensors=True) scores = [ tf.nn.softmax(x) for x in tf.split(scores, len(actions) // 10, axis=0) ] scores = tf.concat(scores, axis=0) rews = rews * tf.expand_dims(scores, axis=-1) grads = self._compute_gradients(actions, rews, weights, **kwargs) grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm) grads_and_vars = zip(grads, self.trainable_variables) new_vars = [v - self.learning_rate * g for g, v in grads_and_vars] self.optimizer.apply_gradients(grads_and_vars) grads_loss = self._compute_gradients( dev_actions, dev_rews, dev_weights, loss_str='dev', use_entropy_regularization=False, **dev_kwargs) score_grads = tape0.gradient( new_vars, self._score_vars, output_gradients=grads_loss) del tape0 score_grads_and_vars = self._score_grad_clipping( zip(score_grads, self._score_vars)) self.score_optimizer.apply_gradients( score_grads_and_vars, global_step=self.global_step) if self.log_summaries: grads = list(zip(*grads_and_vars)[0]) score_grads = list(zip(*score_grads_and_vars)[0]) contrib_summary.scalar('global_norm/train_grad', tf.global_norm(grads)) contrib_summary.scalar('global_norm/meta_grad', tf.global_norm(score_grads)) if self._debug and (self._counter % self.log_every == 0): tf.print( 'Epoch {} scores='.format(self._counter), scores[:20], summarize=10, output_stream=sys.stdout) self._counter += 1
def get_product_scores(model, user_idxs, query_word_idx, product_idxs=None, scope=None): with variable_scope.variable_scope(scope or "embedding_graph"): # get user embedding [None, embed_size] user_vec = tf.nn.embedding_lookup(model.user_emb, user_idxs) # get query embedding [None, embed_size] query_vec, query_embs = get_query_embedding(model, query_word_idx, True) model.print_ops.append(tf.print("user_vec", user_vec)) model.print_ops.append(tf.print("query_vec", query_vec)) # get candidate product embedding [None, embed_size] product_vec = None product_bias = None if product_idxs != None: product_vec = tf.nn.embedding_lookup(model.product_emb, product_idxs) product_bias = tf.nn.embedding_lookup(model.product_bias, product_idxs) else: product_vec = model.product_emb product_bias = model.product_bias print('Similarity Function : ' + model.similarity_func) model.print_ops.append(tf.print("product vec", product_vec)) if model.similarity_func == 'product': return tf.matmul( (1.0 - model.Wu) * user_vec + model.Wu * query_vec, product_vec, transpose_b=True) elif model.similarity_func == 'bias_product': return tf.matmul( (1.0 - model.Wu) * user_vec + model.Wu * query_vec, product_vec, transpose_b=True) + product_bias else: user_vec = user_vec / tf.sqrt( tf.reduce_sum(tf.square(user_vec), 1, keep_dims=True)) query_vec = query_vec / tf.sqrt( tf.reduce_sum(tf.square(query_vec), 1, keep_dims=True)) product_vec = product_vec / tf.sqrt( tf.reduce_sum(tf.square(product_vec), 1, keep_dims=True)) return tf.matmul( (1.0 - model.Wu) * user_vec + model.Wu * query_vec, product_vec, transpose_b=True)
def log_prob_fn(params): rho, alpha, sigma = tf.split(params, [num_features, 1, 1], -1) one = tf.ones(num_features) def indep(d): return tfd.Independent(d, 1) p_rho = indep(tfd.InverseGamma(5. * one, 5. * one)) p_alpha = indep(tfd.HalfNormal([1.])) p_sigma = indep(tfd.HalfNormal([1.])) rho_shape = tf.shape(rho) alpha_shape = tf.shape(alpha) x1 = tf.expand_dims(x, -2) x2 = tf.expand_dims(x, -3) exp = -0.5 * tf.squared_difference(x1, x2) exp /= tf.reshape(tf.square(rho), tf.concat([rho_shape[:1], [1, 1], rho_shape[1:]], 0)) exp = tf.reduce_sum(exp, -1, keep_dims=True) exp += 2. * tf.reshape(tf.log(alpha), tf.concat([alpha_shape[:1], [1, 1], alpha_shape[1:]], 0)) exp = tf.exp(exp[Ellipsis, 0]) exp += tf.matrix_diag(tf.tile(tf.square(sigma), [1, int(x.shape[0])]) + 1e-6) exp = tf.check_numerics(exp, "exp 2 has NaNs") with tf.control_dependencies([tf.print(exp[0], summarize=99999)]): exp = tf.identity(exp) p_y = tfd.MultivariateNormalFullCovariance( covariance_matrix=exp) log_prob = ( p_rho.log_prob(rho) + p_alpha.log_prob(alpha) + p_sigma.log_prob(sigma) + p_y.log_prob(y)) return log_prob
def body(spec_idx, attack, success): """Runs a separate PGD attack for each specification.""" adversarial_input = pgd_attack( build_loss_fn(spec_idx), duplicated_inputs, epsilon=self._epsilon, num_steps=self._num_steps, image_bounds=self._input_bounds, random_init=self._random_init, optimizer=optimizer, project_perturbation=self._project_perturbation) new_attack = self.find_worst_attack(flat_objective_fn, adversarial_input, batch_size, input_shape) new_logits = self._eval_fn(new_attack) # Count the number of sample that violate any specification. new_success = _any_greater( self._specification.evaluate(new_logits)) # The first iteration always sets the attack and logits. use_new_values = tf.logical_or(tf.equal(spec_idx, 0), new_success) print_op = tf.print('Processed specification #', spec_idx) with tf.control_dependencies([print_op]): new_spec_idx = spec_idx + 1 return (new_spec_idx, tf.where(use_new_values, new_attack, attack), tf.logical_or(success, new_success))
def _apply_reapated_text_masking( config: RetrieverConfig, question_hash: tf.Tensor, question_hash_transposed: tf.Tensor, labels: tf.Tensor, logits: tf.Tensor, ) -> tf.Tensor: """Applies repated text masking. Args: config: Retriever config. question_hash: <int64>[global_batch_size, 1] question_hash_transposed: <int64>[1, batch_size] labels: <int64>[batch_size, global_batch_size * num_tables] logits: <float>[batch_size, global_batch_size * num_tables] Returns: Masked logits (same shape / dtype). """ # Make sure not all hashes are 0. # This indicates the "question_hash" feature wasn't set. assert_op = tf.assert_equal( tf.math.reduce_all(tf.math.equal(question_hash, 0)), [False]) with tf.control_dependencies([assert_op]): logging.vlog(2, "question_hash: %s", question_hash) logging.vlog(2, "question_hash_transposed: %s", question_hash_transposed) logging.vlog(2, "labels: %s", labels) logging.vlog(2, "logits: %s", logits) # <bool>[batch_size, global_batch_size] repeated_texts = tf.math.equal(question_hash, question_hash_transposed) if config.use_mined_negatives: batch_size = repeated_texts.shape[0] global_batch_size = repeated_texts.shape[1] num_tables = logits.shape[1] // global_batch_size # <bool>[batch_size, global_batch_size * num_tables] repeated_texts = tf.concat([ repeated_texts, tf.zeros(shape=(batch_size, (num_tables - 1) * global_batch_size), dtype=tf.bool) ], axis=1) repeated_texts = ( repeated_texts # Makes sure original correct question pair isn't masked & tf.math.equal(labels, 0)) logging.vlog(2, "repeated texts: %s", repeated_texts) ops = [] if logging.vlog_is_on(2): ops.append( tf.print( "repeated texts content:", question_hash, repeated_texts, output_stream=logging.info, )) with tf.control_dependencies(ops): return tf.where(repeated_texts, tf.zeros_like(logits) - _INF, logits)
def update(self, samples, contexts): """Update the policy based on the training samples and contexts.""" # To prevent memory leaks in tf eager if self._counter % self.log_every == 0: tf.set_random_seed(self._seed) batch_actions, batch_rews, batch_weights, kwargs = \ self.create_batch(samples, contexts=contexts) grads = self._compute_gradients(batch_actions, batch_rews, batch_weights, **kwargs) grads, norm = tf.clip_by_global_norm(grads, self.max_grad_norm) if self._debug and self._counter % self.log_every == 0: tf.print('Epoch {}: Grad norm='.format(self._counter), norm, output_stream=sys.stdout) self.optimizer.apply_gradients(zip(grads, self.trainable_variables), global_step=self.global_step) self._counter += 1
def print_filtered_subset(ex): """Print filtered subset for debug purpose.""" if isinstance(ex, dict) and 'id' in ex and 'label' in ex: print_op = tf.print('filtered_example:', ex['id'], ex['label'], output_stream=tf.logging.error) with tf.control_dependencies([print_op]): ex['id'] = tf.identity(ex['id']) return ex
def create_print_op(): return tf.print(_MLPERF_LOG_PREFIX, self.mlperf_model_name, tf.timestamp(), caller, key, ': { "deferred": true, "value":', tensor_value, '}', output_stream=sys.stdout)
def get_summaries(ops): summaries = [] for name, op in ops.items(): # Ensure to log the value ops before writing them in the summary. # We do this instead of a hook to ensure IS/FID are never computed twice. print_op = tf.print(name, [op], output_stream=tf.logging.info) with tf.control_dependencies([print_op]): summary = tf.summary.scalar(name, op) summaries.append(summary) return summaries
def compute_seq_metrics(label_dict, feature_dict, debug=False, mask=None): """Compute the reference accuracy.""" gt_lengths = verb_refs_to_lengths(label_dict["task"], label_dict["verb_refs"], include_eos=False) pred_lengths = verb_refs_to_lengths(feature_dict["task"], feature_dict["verb_refs"], include_eos=False) gt_actions = tf.concat([ tf.expand_dims(label_dict["verbs"], 2), tf.expand_dims(label_dict["objects"], 2), label_dict["input_refs"] ], axis=-1) pr_actions = tf.concat([ tf.expand_dims(feature_dict["verbs"], 2), tf.expand_dims(feature_dict["objects"], 2), feature_dict["input_refs"] ], axis=-1) complete_act_acc, partial_act_acc = sequence_accuracy(gt_actions, pr_actions, gt_lengths, pred_lengths, debug=debug, name="act") gt_refs = tf.concat([ label_dict["verb_refs"], label_dict["obj_refs"], label_dict["input_refs"] ], axis=-1) pr_refs = tf.concat([ feature_dict["verb_refs"], feature_dict["obj_refs"], feature_dict["input_refs"] ], axis=-1) if mask is not None: mask = tf.expand_dims(tf.expand_dims(mask, 0), 0) gt_refs = gt_refs * mask pr_refs = pr_refs * mask pred_lengths = gt_lengths with tf.control_dependencies( [tf.print("mask", gt_refs, pr_refs, summarize=100)]): complete_refs_acc, partial_refs_acc = sequence_accuracy(gt_refs, pr_refs, gt_lengths, pred_lengths, debug=debug, name="ref") refs_metrics = {} refs_metrics["complete_acts_acc"] = complete_act_acc refs_metrics["partial_acts_acc"] = partial_act_acc refs_metrics["complete_refs_acc"] = complete_refs_acc refs_metrics["partial_refs_acc"] = partial_refs_acc refs_metrics["gt_seq"] = gt_actions refs_metrics["pred_seq"] = pr_actions return refs_metrics
def sequence_accuracy(gt_seqs, decode_seqs, gt_seq_lengths, pr_seq_lengths, debug=False, name=""): """Computes the complete and the partial sequence accuracy.""" gt_shape = common_layers.shape_list(gt_seqs) pr_shape = common_layers.shape_list(decode_seqs) batch_size = gt_shape[0] depth = gt_shape[-1] gt_len = gt_shape[1] pr_len = pr_shape[1] max_len = tf.maximum(gt_len, pr_len) gt_seqs = tf.pad(gt_seqs, [[0, 0], [0, max_len - gt_len], [0, 0]]) decode_seqs = tf.pad(decode_seqs, [[0, 0], [0, max_len - pr_len], [0, 0]]) gt_seqs = tf.where( tf.tile( tf.expand_dims(tf.sequence_mask(gt_seq_lengths, maxlen=max_len), 2), [1, 1, depth]), gt_seqs, tf.fill(tf.shape(gt_seqs), -1)) decode_seqs = tf.where( tf.tile( tf.expand_dims(tf.sequence_mask(pr_seq_lengths, maxlen=max_len), 2), [1, 1, depth]), decode_seqs, tf.fill(tf.shape(decode_seqs), -1)) # [batch_size, decode_length] corrects = tf.reduce_all(tf.equal(gt_seqs, decode_seqs), -1) correct_mask = tf.reduce_all(corrects, -1) # [batch_size] if debug: incorrect_mask = tf.logical_not(correct_mask) incorrect_gt = tf.boolean_mask(gt_seqs, incorrect_mask) incorrect_pr = tf.boolean_mask(decode_seqs, incorrect_mask) with tf.control_dependencies([ tf.print(name + "_mismatch", incorrect_gt, incorrect_pr, summarize=1000) ]): correct_mask = tf.identity(correct_mask) correct_seqs = tf.to_float(correct_mask) total_correct_seqs = tf.reduce_sum(correct_seqs) mean_complete_accuracy = total_correct_seqs / tf.to_float(batch_size) # Compute partial accuracy errors = tf.logical_not(corrects) errors = tf.cast(tf.cumsum(tf.to_float(errors), axis=-1), tf.bool) # [batch_size] correct_steps = tf.reduce_sum(tf.to_float(tf.logical_not(errors)), axis=-1) mean_partial_accuracy = tf.reduce_mean( tf.div(tf.minimum(correct_steps, gt_seq_lengths), gt_seq_lengths)) return mean_complete_accuracy, mean_partial_accuracy
def DecayJudge(): # calculate mean new_loss_last = loss_last.assign(loss_mean, use_locking=True) with tf.control_dependencies([new_loss_last]): new_loss_mean = loss_mean.assign(loss_sum / loss_count, use_locking=True) with tf.control_dependencies([new_loss_mean]): clear_sum = loss_sum.assign(0.0, use_locking=True) clear_count = loss_count.assign(0.0, use_locking=True) with tf.control_dependencies([clear_sum, clear_count]): cal_mean = tf.no_op('CalculateMean') # decay judge with tf.control_dependencies([cal_mean]): print_op = tf.print('last loss mean: ', new_loss_last, '\nloss mean: ', new_loss_mean) with tf.control_dependencies([print_op]): decay_on = new_loss_mean * (1 + min_delta) > new_loss_last new_lr_mul = tf.cond( decay_on, lambda: tf.math.maximum(min_mul, lr_mul * factor), lambda: lr_mul) print_op = tf.print('new LR multiplier: ', new_lr_mul) with tf.control_dependencies([print_op]): return lr_mul.assign(new_lr_mul, use_locking=True)
def py_print_iteration_info(msg, var, n, debug=True): """adds a tf.print op to the graph while ensuring it will run (when the output is used).""" if not debug: return var var_print = tf.print("n=", n, "\t", msg, tf.shape(var), var, summarize=-1, output_stream=sys.stdout) with tf.control_dependencies([var_print]): var = tf.identity(var) return var
def multiline_print(lists): """Prints multiple lines of output using tf.print.""" combined_list = [] combined_list += lists[0] # We prepend newline characters to strings at the start of lines to avoid # the ugly space intendations that tf.print's behavior of separating # everything with a space would otherwise cause. for item in lists[1:]: if isinstance(item[0], str): combined_list += (("\n" + item[0],) + item[1:]) else: combined_list += (("\n",) + item) return tf.print(*combined_list)
def testLoss(self): batch_size = 2 key_depth = 5 val_depth = 5 memory_size = 4 window_size = 3 x_depth = 5 memory = transformer_memory.TransformerMemory(batch_size, key_depth, val_depth, memory_size) x = tf.random_uniform([batch_size, window_size, x_depth], minval=.0) memory_results, _, _, _ = (memory.pre_attention( tf.random_uniform([batch_size], minval=0, maxval=1, dtype=tf.int32), x, None, None)) x = memory.post_attention(memory_results, x) with tf.control_dependencies([tf.print("x", x)]): is_nan = tf.reduce_any(tf.math.is_nan(x)) with self.test_session() as session: session.run(tf.global_variables_initializer()) for _ in range(100): is_nan_value, _ = session.run([is_nan, x]) self.assertEqual(is_nan_value, False)
def train( root_dir, load_root_dir=None, env_load_fn=None, env_name=None, num_parallel_environments=1, # pylint: disable=unused-argument agent_class=None, initial_collect_random=True, # pylint: disable=unused-argument initial_collect_driver_class=None, collect_driver_class=None, num_global_steps=1000000, train_steps_per_iteration=1, train_metrics=None, # Safety Critic training args train_sc_steps=10, train_sc_interval=300, online_critic=False, # Params for eval run_eval=False, num_eval_episodes=30, eval_interval=1000, eval_metrics_callback=None, # Params for summaries and logging train_checkpoint_interval=10000, policy_checkpoint_interval=5000, rb_checkpoint_interval=20000, keep_rb_checkpoint=False, log_interval=1000, summary_interval=1000, summaries_flush_secs=10, early_termination_fn=None, env_metric_factories=None): # pylint: disable=unused-argument """A simple train and eval for SC-SAC.""" root_dir = os.path.expanduser(root_dir) train_dir = os.path.join(root_dir, 'train') train_summary_writer = tf.compat.v2.summary.create_file_writer( train_dir, flush_millis=summaries_flush_secs * 1000) train_summary_writer.set_as_default() train_metrics = train_metrics or [] if run_eval: eval_dir = os.path.join(root_dir, 'eval') eval_summary_writer = tf.compat.v2.summary.create_file_writer( eval_dir, flush_millis=summaries_flush_secs * 1000) eval_metrics = [ tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes), tf_metrics.AverageEpisodeLengthMetric( buffer_size=num_eval_episodes), ] + [tf_py_metric.TFPyMetric(m) for m in train_metrics] global_step = tf.compat.v1.train.get_or_create_global_step() with tf.compat.v2.summary.record_if( lambda: tf.math.equal(global_step % summary_interval, 0)): tf_env = env_load_fn(env_name) if not isinstance(tf_env, tf_py_environment.TFPyEnvironment): tf_env = tf_py_environment.TFPyEnvironment(tf_env) if run_eval: eval_py_env = env_load_fn(env_name) eval_tf_env = tf_py_environment.TFPyEnvironment(eval_py_env) time_step_spec = tf_env.time_step_spec() observation_spec = time_step_spec.observation action_spec = tf_env.action_spec() print('obs spec:', observation_spec) print('action spec:', action_spec) if online_critic: resample_metric = tf_py_metric.TfPyMetric( py_metrics.CounterMetric('unsafe_ac_samples')) tf_agent = agent_class(time_step_spec, action_spec, train_step_counter=global_step, resample_metric=resample_metric) else: tf_agent = agent_class(time_step_spec, action_spec, train_step_counter=global_step) tf_agent.initialize() # Make the replay buffer. collect_data_spec = tf_agent.collect_data_spec logging.info('Allocating replay buffer ...') # Add to replay buffer and other agent specific observers. replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer( collect_data_spec, max_length=1000000) logging.info('RB capacity: %i', replay_buffer.capacity) logging.info('ReplayBuffer Collect data spec: %s', collect_data_spec) agent_observers = [replay_buffer.add_batch] if online_critic: online_replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer( collect_data_spec, max_length=10000) online_rb_ckpt_dir = os.path.join(train_dir, 'online_replay_buffer') online_rb_checkpointer = common.Checkpointer( ckpt_dir=online_rb_ckpt_dir, max_to_keep=1, replay_buffer=online_replay_buffer) clear_rb = common.function(online_replay_buffer.clear) agent_observers.append(online_replay_buffer.add_batch) train_metrics = [ tf_metrics.NumberOfEpisodes(), tf_metrics.EnvironmentSteps(), tf_metrics.AverageReturnMetric(buffer_size=num_eval_episodes, batch_size=tf_env.batch_size), tf_metrics.AverageEpisodeLengthMetric( buffer_size=num_eval_episodes, batch_size=tf_env.batch_size), ] + [tf_py_metric.TFPyMetric(m) for m in train_metrics] if not online_critic: eval_policy = tf_agent.policy else: eval_policy = tf_agent._safe_policy # pylint: disable=protected-access initial_collect_policy = random_tf_policy.RandomTFPolicy( time_step_spec, action_spec) if not online_critic: collect_policy = tf_agent.collect_policy else: collect_policy = tf_agent._safe_policy # pylint: disable=protected-access train_checkpointer = common.Checkpointer( ckpt_dir=train_dir, agent=tf_agent, global_step=global_step, metrics=metric_utils.MetricsGroup(train_metrics, 'train_metrics')) policy_checkpointer = common.Checkpointer(ckpt_dir=os.path.join( train_dir, 'policy'), policy=eval_policy, global_step=global_step) safety_critic_checkpointer = common.Checkpointer( ckpt_dir=os.path.join(train_dir, 'safety_critic'), safety_critic=tf_agent._safety_critic_network, # pylint: disable=protected-access global_step=global_step) rb_ckpt_dir = os.path.join(train_dir, 'replay_buffer') rb_checkpointer = common.Checkpointer(ckpt_dir=rb_ckpt_dir, max_to_keep=1, replay_buffer=replay_buffer) if load_root_dir: load_root_dir = os.path.expanduser(load_root_dir) load_train_dir = os.path.join(load_root_dir, 'train') misc.load_pi_ckpt(load_train_dir, tf_agent) # loads tf_agent if load_root_dir is None: train_checkpointer.initialize_or_restore() rb_checkpointer.initialize_or_restore() safety_critic_checkpointer.initialize_or_restore() collect_driver = collect_driver_class(tf_env, collect_policy, observers=agent_observers + train_metrics) collect_driver.run = common.function(collect_driver.run) tf_agent.train = common.function(tf_agent.train) if not rb_checkpointer.checkpoint_exists: logging.info('Performing initial collection ...') common.function( initial_collect_driver_class(tf_env, initial_collect_policy, observers=agent_observers + train_metrics).run)() last_id = replay_buffer._get_last_id() # pylint: disable=protected-access logging.info('Data saved after initial collection: %d steps', last_id) tf.print( replay_buffer._get_rows_for_id(last_id), # pylint: disable=protected-access output_stream=logging.info) if run_eval: results = metric_utils.eager_compute( eval_metrics, eval_tf_env, eval_policy, num_episodes=num_eval_episodes, train_step=global_step, summary_writer=eval_summary_writer, summary_prefix='Metrics', ) if eval_metrics_callback is not None: eval_metrics_callback(results, global_step.numpy()) metric_utils.log_metrics(eval_metrics) if FLAGS.viz_pm: eval_fig_dir = osp.join(eval_dir, 'figs') if not tf.io.gfile.isdir(eval_fig_dir): tf.io.gfile.makedirs(eval_fig_dir) time_step = None policy_state = collect_policy.get_initial_state(tf_env.batch_size) timed_at_step = global_step.numpy() time_acc = 0 # Dataset generates trajectories with shape [Bx2x...] dataset = replay_buffer.as_dataset(num_parallel_calls=3, num_steps=2).prefetch(3) iterator = iter(dataset) if online_critic: online_dataset = online_replay_buffer.as_dataset( num_parallel_calls=3, num_steps=2).prefetch(3) online_iterator = iter(online_dataset) @common.function def critic_train_step(): """Builds critic training step.""" experience, buf_info = next(online_iterator) if env_name in [ 'IndianWell', 'IndianWell2', 'IndianWell3', 'DrunkSpider', 'DrunkSpiderShort' ]: safe_rew = experience.observation['task_agn_rew'] else: safe_rew = agents.process_replay_buffer( online_replay_buffer, as_tensor=True) safe_rew = tf.gather(safe_rew, tf.squeeze(buf_info.ids), axis=1) ret = tf_agent.train_sc(experience, safe_rew) clear_rb() return ret @common.function def train_step(): experience, _ = next(iterator) ret = tf_agent.train(experience) return ret if not early_termination_fn: early_termination_fn = lambda: False loss_diverged = False # How many consecutive steps was loss diverged for. loss_divergence_counter = 0 mean_train_loss = tf.keras.metrics.Mean(name='mean_train_loss') if online_critic: mean_resample_ac = tf.keras.metrics.Mean( name='mean_unsafe_ac_samples') resample_metric.reset() while (global_step.numpy() <= num_global_steps and not early_termination_fn()): # Collect and train. start_time = time.time() time_step, policy_state = collect_driver.run( time_step=time_step, policy_state=policy_state, ) if online_critic: mean_resample_ac(resample_metric.result()) resample_metric.reset() if time_step.is_last(): resample_ac_freq = mean_resample_ac.result() mean_resample_ac.reset_states() tf.compat.v2.summary.scalar(name='unsafe_ac_samples', data=resample_ac_freq, step=global_step) for _ in range(train_steps_per_iteration): train_loss = train_step() mean_train_loss(train_loss.loss) if online_critic: if global_step.numpy() % train_sc_interval == 0: for _ in range(train_sc_steps): sc_loss, lambda_loss = critic_train_step() # pylint: disable=unused-variable total_loss = mean_train_loss.result() mean_train_loss.reset_states() # Check for exploding losses. if (math.isnan(total_loss) or math.isinf(total_loss) or total_loss > MAX_LOSS): loss_divergence_counter += 1 if loss_divergence_counter > TERMINATE_AFTER_DIVERGED_LOSS_STEPS: loss_diverged = True break else: loss_divergence_counter = 0 time_acc += time.time() - start_time if global_step.numpy() % log_interval == 0: logging.info('step = %d, loss = %f', global_step.numpy(), total_loss) steps_per_sec = (global_step.numpy() - timed_at_step) / time_acc logging.info('%.3f steps/sec', steps_per_sec) tf.compat.v2.summary.scalar(name='global_steps_per_sec', data=steps_per_sec, step=global_step) timed_at_step = global_step.numpy() time_acc = 0 for train_metric in train_metrics: train_metric.tf_summaries(train_step=global_step, step_metrics=train_metrics[:2]) global_step_val = global_step.numpy() if global_step_val % train_checkpoint_interval == 0: train_checkpointer.save(global_step=global_step_val) if global_step_val % policy_checkpoint_interval == 0: policy_checkpointer.save(global_step=global_step_val) safety_critic_checkpointer.save(global_step=global_step_val) if global_step_val % rb_checkpoint_interval == 0: if online_critic: online_rb_checkpointer.save(global_step=global_step_val) rb_checkpointer.save(global_step=global_step_val) if run_eval and global_step.numpy() % eval_interval == 0: results = metric_utils.eager_compute( eval_metrics, eval_tf_env, eval_policy, num_episodes=num_eval_episodes, train_step=global_step, summary_writer=eval_summary_writer, summary_prefix='Metrics', ) if eval_metrics_callback is not None: eval_metrics_callback(results, global_step.numpy()) metric_utils.log_metrics(eval_metrics) if FLAGS.viz_pm: savepath = 'step{}.png'.format(global_step_val) savepath = osp.join(eval_fig_dir, savepath) misc.record_episode_vis_summary(eval_tf_env, eval_policy, savepath) if not keep_rb_checkpoint: misc.cleanup_checkpoints(rb_ckpt_dir) if loss_diverged: # Raise an error at the very end after the cleanup. raise ValueError('Loss diverged to {} at step {}, terminating.'.format( total_loss, global_step.numpy())) return total_loss
import tensorflow.compat.v1 as tf tf.disable_v2_behavior() hello = tf.constant("Hello TF!") sess = tf.Session() print(sess.run(hello)) node1 = tf.constant(3.0, tf.float32) node2 = tf.constant(4.0) node3 = tf.add(node1, node2) tf.print(f"node1 : {node1}, node2 : {node2}") tf.print(f"node3 : {node3}") print(sess.run(node3)) ph1 = tf.placeholder(tf.float32) ph2 = tf.placeholder(tf.float32) add = ph1 + ph2 @tf.function def forward(): return ph1 + ph2 print(sess.run(add, feed_dict={ph1: 3, ph2: 4.5})) print(sess.run(add, feed_dict={ph1: [1, 3], ph2: [2, 4]}))
def __init__(self, train_batch_size=4096, test_chain_batch_size=4096, bijector="iaf", log_dir="/tmp/neutra", base_learning_rate=1e-3, q_base_scale=1., learning_rate_schedule=[[6000, 1e-1]]): target, target_spec = GetTargetSpec() self.target = target self.target_spec = target_spec with gin.config_scope("train"): train_target, train_target_spec = GetTargetSpec() self.train_target = train_target self.train_target_spec = train_target_spec if bijector == "rnvp": bijector_fn = tf.make_template("bijector", MakeRNVPBijectorFn, num_dims=self.target_spec.num_dims) elif bijector == "iaf": bijector_fn = tf.make_template("bijector", MakeIAFBijectorFn, num_dims=self.target_spec.num_dims) elif bijector == "affine": bijector_fn = tf.make_template("bijector", MakeAffineBijectorFn, num_dims=self.target_spec.num_dims) else: bijector_fn = lambda *args, **kwargs: tfb.Identity() self.train_bijector = bijector_fn(train=True) self.bijector = bijector_fn(train=False) if train_target_spec.bijector is not None: print("Using train target bijector") self.train_bijector = tfb.Chain( [train_target_spec.bijector, self.train_bijector]) if target_spec.bijector is not None: print("Using target bijector") self.bijector = tfb.Chain([target_spec.bijector, self.bijector]) q_base = tfd.Independent( tfd.Normal(loc=tf.zeros(self.target_spec.num_dims), scale=q_base_scale * tf.ones(self.target_spec.num_dims)), 1) self.q_x_train = tfd.TransformedDistribution(q_base, self.train_bijector) self.q_x = tfd.TransformedDistribution(q_base, self.bijector) # Params self.train_batch_size = int(train_batch_size) self.test_chain_batch_size = tf.placeholder_with_default( test_chain_batch_size, [], "test_chain_batch_size") self.test_batch_size = tf.placeholder_with_default( 16384 * 8, [], "test_batch_size") self.test_num_steps = tf.placeholder_with_default( 1000, [], "test_num_steps") self.test_num_leapfrog_steps = tf.placeholder_with_default( tf.to_int32(2), [], "test_num_leapfrog_steps") self.test_step_size = tf.placeholder_with_default( 0.1, [], "test_step_size") # Test self.neutra_outputs = MakeNeuTra( target=self.target, q=self.q_x, batch_size=self.test_chain_batch_size, num_steps=self.test_num_steps, num_leapfrog_steps=self.test_num_leapfrog_steps, step_size=self.test_step_size, ) self.z_chain = tf.reshape( self.bijector.inverse( tf.reshape(self.neutra_outputs.x_chain, [-1, self.target_spec.num_dims])), tf.shape(self.neutra_outputs.x_chain)) self.target_samples = self.target.sample(self.test_batch_size) self.target_z = self.bijector.inverse(self.target_samples) self.q_samples = self.q_x.sample(self.test_batch_size) self.target_cov = utils.Covariance(self.target_samples) self.target_eigvals, self.target_eigvecs = tf.linalg.eigh( self.target_cov) self.cached_target_eigvals = tf.get_local_variable( "cached_target_eigvals", self.target_eigvals.shape, initializer=tf.zeros_initializer()) self.cached_target_eigvecs = tf.get_local_variable( "cached_target_eigvecs", self.target_eigvecs.shape, initializer=tf.zeros_initializer()) self.cached_target_stats_update_op = [ self.cached_target_eigvals.assign(self.target_eigvals), self.cached_target_eigvecs.assign(self.target_eigvecs), tf.print("Assigning target stats") ] def variance(x): x -= tf.reduce_mean(x, 0, keep_dims=True) x = tf.square(x) return x def rotated_variance(x): x2 = tf.reshape(x, [-1, self.target_spec.num_dims]) x2 -= tf.reduce_mean(x2, 0, keep_dims=True) x2 = tf.matmul(x2, self.cached_target_eigvecs) x2 = tf.square(x2) return tf.reshape(x2, tf.shape(x)) functions = [ ("mean", tf.identity), # ("var", variance), ("square", tf.square), # ("rot_square", rot_square), # ("rot_var", rotated_variance), ] self.cached_target_mean = {} self.cached_target_mean_update_op = [ tf.print("Assigning target means.") ] self.neutra_stats = {} self.q_stats = {} for name, f in functions: target_mean = tf.reduce_mean(f(self.target_samples), 0) cached_target_mean = tf.get_local_variable(name + "_cached_mean", target_mean.shape) if self.target_spec.stats is not None: self.cached_target_mean_update_op.append( cached_target_mean.assign(self.target_spec.stats[name])) else: self.cached_target_mean_update_op.append( cached_target_mean.assign(target_mean)) self.cached_target_mean[name] = cached_target_mean self.q_stats[name] = ComputeQStats(f(self.q_samples), cached_target_mean) self.neutra_stats[name] = ComputeChainStats( f(self.neutra_outputs.x_chain), cached_target_mean, self.test_num_leapfrog_steps) # Training self.train_q_samples = self.q_x_train.sample(self.train_batch_size) self.train_log_q_x = self.q_x_train.log_prob(self.train_q_samples) self.kl_q_p = tf.reduce_mean( self.train_log_q_x - self.target.log_prob(self.train_q_samples)) loss = self.kl_q_p reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES) if reg_losses: tf.logging.info("Regularizing.") loss += tf.add_n(reg_losses) self.loss = tf.check_numerics(loss, "Loss has NaNs") self.global_step = tf.train.get_or_create_global_step() steps, factors = list(zip(*learning_rate_schedule)) learning_rate = base_learning_rate * tf.train.piecewise_constant( self.global_step, steps, [1.0] + list(factors)) opt = tf.train.AdamOptimizer(learning_rate=learning_rate) self.train_op = opt.minimize(self.loss, global_step=self.global_step) tf.summary.scalar("kl_q_p", self.kl_q_p) tf.summary.scalar("loss", self.loss) self.init = [ tf.global_variables_initializer(), tf.local_variables_initializer(), tf.print("Initializing variables") ] self.saver = tf.train.Saver() self.log_dir = log_dir
def decode_fn(hparams): """The main function.""" decode_dict, decode_mask, label_dict = _decode_common(hparams) if FLAGS.problem != "android_howto": decode_dict["input_refs"] = decode_utils.unify_input_ref( decode_dict["verbs"], decode_dict["input_refs"]) print_ops = [] for key in [ "raw_task", "verbs", "objects", "verb_refs", "obj_refs", "input_refs" ]: print_ops.append( tf.print(key, tf.shape(decode_dict[key]), decode_dict[key], label_dict[key], "decode_mask", decode_mask, summarize=100)) acc_metrics = decode_utils.compute_seq_metrics(label_dict, decode_dict, mask=None) saver = tf.train.Saver() with tf.Session() as session: session.run(tf.global_variables_initializer()) latest_checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoint_path) tf.logging.info("Restoring from the latest checkpoint: %s" % (latest_checkpoint)) saver.restore(session, latest_checkpoint) task_seqs = [] ref_seqs = [] act_seqs = [] mask_seqs = [] try: i = 0 while True: tf.logging.info("Example %d" % i) task, acc, mask, label, decode = session.run([ decode_dict["raw_task"], acc_metrics, decode_mask, label_dict, decode_dict ]) ref_seq = {} ref_seq["gt_seq"] = np.concatenate([ label["verb_refs"], label["obj_refs"], label["input_refs"] ], axis=-1) ref_seq["pred_seq"] = np.concatenate([ decode["verb_refs"], decode["obj_refs"], decode["input_refs"] ], axis=-1) ref_seq["complete_seq_acc"] = acc["complete_refs_acc"] ref_seq["partial_seq_acc"] = acc["partial_refs_acc"] act_seq = {} act_seq["gt_seq"] = np.concatenate([ np.expand_dims(label["verbs"], 2), np.expand_dims(label["objects"], 2), label["input_refs"] ], axis=-1) act_seq["pred_seq"] = np.concatenate([ np.expand_dims(decode["verbs"], 2), np.expand_dims(decode["objects"], 2), decode["input_refs"] ], axis=-1) act_seq["complete_seq_acc"] = acc["complete_acts_acc"] act_seq["partial_seq_acc"] = acc["partial_acts_acc"] print("task", task) print("ref_seq", ref_seq) print("act_seq", act_seq) print("mask", mask) task_seqs.append(task) ref_seqs.append(ref_seq) act_seqs.append(act_seq) mask_seqs.append(mask) i += 1 except tf.errors.OutOfRangeError: pass save(task_seqs, ref_seqs, mask_seqs, "joint_refs") save(task_seqs, act_seqs, mask_seqs, "joint_act")
def tf_print(x): tf.print(x) return x + 1
""" import tensorflow.compat.v1 as tf tf.disable_v2_behavior() a = tf.constant(1) b = tf.constant(2) c = tf.add(a,b) print(c) sess = tf.Session() print(sess.run([a, b, c])) tf.print(a,b,c) sess.close() # 2 a = tf.placeholder(tf.float32) b = tf.placeholder(tf.float32) adder_node = a + b sess = tf.Session() print(sess.run(adder_none, feed_dict = {a:3, b:4.5})) print(sess.run(adder_none, feed_dict = {a:[1,3], b: [2,4]}))
def bottleneck(inputs, depth, depth_bottleneck, stride, rate=1, outputs_collections=None, scope=None, use_bounded_activations=False, initializers=None, insert_shift=False, split_model=False): """Bottleneck residual unit variant with BN after convolutions. This is the original residual unit proposed in [1]. See Fig. 1(a) of [2] for its definition. Note that we use here the bottleneck variant which has an extra bottleneck layer. When putting together two consecutive ResNet blocks that use this unit, one should use stride = 2 in the last unit of the first block. Args: inputs: A tensor of size [batch, height, width, channels]. depth: The depth of the ResNet unit output. depth_bottleneck: The depth of the bottleneck layers. stride: The ResNet unit's stride. Determines the amount of downsampling of the units output compared to its input. rate: An integer, rate for atrous convolution. outputs_collections: Collection to add the ResNet unit output. scope: Optional variable_scope. use_bounded_activations: Whether or not to use bounded activations. Bounded activations better lend themselves to quantized inference. initializers <MODIFICATION>: Weight and bias initializers for the included conv blocks insert_shift <MODIFICATION>: If true, inserts shift operation in front of every first conv1x1 within a block. split_model <MODIFICATION>: (overrides "shift") If true, inserts placeholder in front of every first conv1x1 within a block. This allows for insertion of shift outside of TF. Returns: The ResNet unit's output. """ with tf.variable_scope(scope, 'bottleneck_v1', [inputs]) as sc: fake_inputs = None fake_shortcut = None if split_model: shift_input = tf.identity(inputs, 'prev_conv_output') #shortcut_input = tf.identity(shortcut, 'prev_shortcut_output') shortcut_input = tf.identity(inputs, 'prev_shortcut_output') fake_inputs = tf.placeholder(tf.float32, shape=inputs.get_shape(), name='conv_input') #fake_shortcut = tf.placeholder(tf.float32, shape=shortcut.get_shape(), name='shortcut_input') fake_shortcut = tf.placeholder(tf.float32, shape=inputs.get_shape(), name='shortcut_input') elif insert_shift: shift_input = tf.identity(inputs, 'prev_conv_output') #shortcut_input = tf.identity(shortcut, 'prev_shortcut_output') shortcut_input = tf.identity(inputs, 'prev_shortcut_output') fold = inputs.get_shape()[-1] // 8 shifted = tf.Variable(tf.zeros_like(inputs)) initial_print = tf.print("initial: ", shifted) shift_left = tf.assign(shifted[:-1, :, :, :fold], shift_input[1:, :, :, :fold]) shift_right = tf.assign(shifted[1:, :, :, fold:2 * fold], shift_input[:-1, :, :, fold:2 * fold]) shift_copy = tf.assign(shifted[:, :, :, 2 * fold:], shift_input[:, :, :, 2 * fold:]) final_print = tf.print("final: ", shifted) with tf.control_dependencies([shift_left, shift_right, shift_copy]): fake_inputs = tf.identity(shifted, 'shifted_input') fake_shortcut = shortcut_input else: fake_inputs = inputs #fake_shortcut = shortcut fake_shortcut = inputs depth_in = slim.utils.last_dimension(inputs.get_shape(), min_rank=4) if depth == depth_in: #shortcut = resnet_utils.subsample(inputs, stride, 'shortcut') shortcut = resnet_utils.subsample(fake_shortcut, stride, 'shortcut') else: shortcut = slim.conv2d( #inputs, fake_shortcut, depth, [1, 1], stride=stride, activation_fn=tf.nn.relu6 if use_bounded_activations else None, scope='shortcut', **initializers["shortcut"]) residual = slim.conv2d(fake_inputs, depth_bottleneck, [1, 1], stride=1, scope='conv1', **initializers["conv1"]) residual = resnet_utils.conv2d_same(residual, depth_bottleneck, 3, stride, rate=rate, scope='conv2', **initializers["conv2"]) residual = slim.conv2d(residual, depth, [1, 1], stride=1, activation_fn=None, scope='conv3', **initializers["conv3"]) if use_bounded_activations: # Use clip_by_value to simulate bandpass activation. residual = tf.clip_by_value(residual, -6.0, 6.0) #output = tf.nn.relu6(fake_shortcut + residual) output = tf.nn.relu6(shortcut + residual) else: #output = tf.nn.relu(fake_shortcut + residual) output = tf.nn.relu(shortcut + residual) return slim.utils.collect_named_outputs(outputs_collections, sc.name, output)
def model_fn(features, labels, mode, params): """The `model_fn` for TPUEstimator.""" del labels, params # Not used. tf.logging.info("*** Features ***") for name in sorted(features.keys()): tf.logging.info(" name = %s, shape = %s", name, features[name].shape) is_training = (mode == tf.estimator.ModeKeys.TRAIN) entity_ids = search_utils.load_database( "entity_ids", [qa_config.num_entities, qa_config.max_entity_len], entity_id_checkpoint, dtype=tf.int32) entity_mask = search_utils.load_database( "entity_mask", [qa_config.num_entities, qa_config.max_entity_len], entity_mask_checkpoint) if FLAGS.model_type == "drkit": # Initialize sparse tensor of ent2ment. with tf.device("/cpu:0"): tf_e2m_data, tf_e2m_indices, tf_e2m_rowsplits = ( search_utils.load_ragged_matrix("ent2ment", e2m_checkpoint)) with tf.name_scope("RaggedConstruction_e2m"): e2m_ragged_ind = tf.RaggedTensor.from_row_splits( values=tf_e2m_indices, row_splits=tf_e2m_rowsplits, validate=False) e2m_ragged_val = tf.RaggedTensor.from_row_splits( values=tf_e2m_data, row_splits=tf_e2m_rowsplits, validate=False) tf_m2e_map = search_utils.load_database("coref", [mips_config.num_mentions], m2e_checkpoint, dtype=tf.int32) total_loss, predictions = create_model_fn( bert_config=bert_config, qa_config=qa_config, mips_config=mips_config, is_training=is_training, features=features, ent2ment_ind=e2m_ragged_ind, ent2ment_val=e2m_ragged_val, ment2ent_map=tf_m2e_map, entity_ids=entity_ids, entity_mask=entity_mask, use_one_hot_embeddings=use_one_hot_embeddings, summary_obj=summary_obj, num_preds=FLAGS.num_preds, is_excluding=FLAGS.is_excluding, ) elif FLAGS.model_type == "drfact": # Initialize sparse tensor of ent2fact. with tf.device("/cpu:0"): # Note: cpu or gpu? tf_e2f_data, tf_e2f_indices, tf_e2f_rowsplits = ( search_utils.load_ragged_matrix("ent2fact", e2f_checkpoint)) with tf.name_scope("RaggedConstruction_e2f"): e2f_ragged_ind = tf.RaggedTensor.from_row_splits( values=tf_e2f_indices, row_splits=tf_e2f_rowsplits, validate=False) e2f_ragged_val = tf.RaggedTensor.from_row_splits( values=tf_e2f_data, row_splits=tf_e2f_rowsplits, validate=False) # Initialize sparse tensor of fact2ent. with tf.device("/cpu:0"): tf_f2e_data, tf_f2e_indices, tf_f2e_rowsplits = ( search_utils.load_ragged_matrix("fact2ent", f2e_checkpoint)) with tf.name_scope("RaggedConstruction_f2e"): f2e_ragged_ind = tf.RaggedTensor.from_row_splits( values=tf_f2e_indices, row_splits=tf_f2e_rowsplits, validate=False) f2e_ragged_val = tf.RaggedTensor.from_row_splits( values=tf_f2e_data, row_splits=tf_f2e_rowsplits, validate=False) # Initialize sparse tensor of fact2fact. with tf.device("/cpu:0"): tf_f2f_data, tf_f2f_indices, tf_f2f_rowsplits = ( search_utils.load_ragged_matrix("fact2fact", f2f_checkpoint)) with tf.name_scope("RaggedConstruction_f2f"): f2f_ragged_ind = tf.RaggedTensor.from_row_splits( values=tf_f2f_indices, row_splits=tf_f2f_rowsplits, validate=False) f2f_ragged_val = tf.RaggedTensor.from_row_splits( values=tf_f2f_data, row_splits=tf_f2f_rowsplits, validate=False) total_loss, predictions = create_model_fn( bert_config=bert_config, qa_config=qa_config, fact_mips_config=fact_mips_config, is_training=is_training, features=features, ent2fact_ind=e2f_ragged_ind, ent2fact_val=e2f_ragged_val, fact2ent_ind=f2e_ragged_ind, fact2ent_val=f2e_ragged_val, fact2fact_ind=f2f_ragged_ind, fact2fact_val=f2f_ragged_val, entity_ids=entity_ids, entity_mask=entity_mask, use_one_hot_embeddings=use_one_hot_embeddings, summary_obj=summary_obj, num_preds=FLAGS.num_preds, is_excluding=FLAGS.is_excluding, ) tvars = tf.trainable_variables() initialized_variable_names = {} scaffold_fn = None if init_checkpoint: (assignment_map, initialized_variable_names) = get_assignment_map_from_checkpoint( tvars, init_checkpoint, load_only_bert=qa_config.load_only_bert) if use_tpu: def tpu_scaffold(): tf.train.init_from_checkpoint(init_checkpoint, assignment_map) return tf.train.Scaffold() scaffold_fn = tpu_scaffold else: tf.train.init_from_checkpoint(init_checkpoint, assignment_map) tf.logging.info("**** Trainable Variables ****") for var in tvars: init_string = "" if var.name in initialized_variable_names: init_string = ", *INIT_FROM_CKPT*" tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, init_string) output_spec = None if mode == tf.estimator.ModeKeys.TRAIN: one_mb = tf.constant(1024 * 1024, dtype=tf.int64) devices = tf.config.experimental.list_logical_devices("GPU") memory_footprints = [] for device in devices: memory_footprint = tf.print( device.name, contrib_memory_stats.MaxBytesInUse() / one_mb, " / ", contrib_memory_stats.BytesLimit() / one_mb) memory_footprints.append(memory_footprint) with tf.control_dependencies(memory_footprints): train_op = create_optimizer(total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu, False) output_spec = tf.estimator.tpu.TPUEstimatorSpec( mode=mode, loss=total_loss, train_op=train_op, scaffold_fn=scaffold_fn) elif mode == tf.estimator.ModeKeys.PREDICT: output_spec = tf.estimator.tpu.TPUEstimatorSpec( mode=mode, predictions=predictions, scaffold_fn=scaffold_fn) else: raise ValueError("Only TRAIN and PREDICT modes are supported: %s" % (mode)) return output_spec
def _eval(metrics, pred_dict, loss_dict, features, areas, compute_seq_accuracy, hparams, metric_types, decode_length=20): """Internal eval function.""" # Assume data sources are not mixed within each batch if compute_seq_accuracy: decode_features = {} for key in features: if not key.endswith("_refs"): decode_features[key] = features[key] decode_utils.decode_n_step(seq2act_model.compute_logits, decode_features, areas, hparams, n=decode_length, beam_size=1) decode_features["input_refs"] = decode_utils.unify_input_ref( decode_features["verbs"], decode_features["input_refs"]) acc_metrics = decode_utils.compute_seq_metrics( features, decode_features) metrics["seq_full_acc"] = tf.metrics.mean(acc_metrics["complete_refs_acc"]) metrics["seq_partial_acc"] = tf.metrics.mean( acc_metrics["partial_refs_acc"]) if "final_accuracy" in metric_types: metrics["complet_act_accuracy"] = tf.metrics.mean( acc_metrics["complete_acts_acc"]) metrics["partial_seq_acc"] = tf.metrics.mean( acc_metrics["partial_acts_acc"]) print0 = tf.print("*** lang", features["raw_task"], summarize=100) with tf.control_dependencies([print0]): loss_dict["total_loss"] = tf.identity(loss_dict["total_loss"]) else: decode_features = None if "ref_accuracy" in metric_types: with tf.control_dependencies([ tf.assert_equal(tf.rank(features["verb_refs"]), 3), tf.assert_equal(tf.shape(features["verb_refs"])[-1], 2)]): _ref_accuracy(features, pred_dict, tf.less(features["verb_refs"][:, :, 0], features["verb_refs"][:, :, 1]), "verb_refs", metrics, decode_features, measure_beginning_eos=True) _ref_accuracy(features, pred_dict, tf.less(features["obj_refs"][:, :, 0], features["obj_refs"][:, :, 1]), "obj_refs", metrics, decode_features, measure_beginning_eos=True) _ref_accuracy(features, pred_dict, tf.less(features["input_refs"][:, :, 0], features["input_refs"][:, :, 1]), "input_refs", metrics, decode_features, measure_beginning_eos=True) if "basic_accuracy" in metric_types: target_verbs = tf.reshape(features["verbs"], [-1]) verb_nonpadding = tf.greater(target_verbs, 1) target_verbs = tf.boolean_mask(target_verbs, verb_nonpadding) predict_verbs = tf.boolean_mask(tf.reshape(pred_dict["verbs"], [-1]), verb_nonpadding) metrics["verb"] = tf.metrics.accuracy( labels=target_verbs, predictions=predict_verbs, name="verb_accuracy") input_mask = tf.reshape( tf.less(features["verb_refs"][:, :, 0], features["verb_refs"][:, :, 1]), [-1]) metrics["input"] = tf.metrics.accuracy( labels=tf.boolean_mask( tf.reshape(tf.to_int32( tf.less(features["input_refs"][:, :, 0], features["input_refs"][:, :, 1])), [-1]), input_mask), predictions=tf.boolean_mask( tf.reshape(pred_dict["input"], [-1]), input_mask), name="input_accuracy") metrics["object"] = tf.metrics.accuracy( labels=tf.boolean_mask(tf.reshape(features["objects"], [-1]), verb_nonpadding), predictions=tf.boolean_mask(tf.reshape(pred_dict["objects"], [-1]), verb_nonpadding), name="object_accuracy") metrics["eval_object_loss"] = tf.metrics.mean( tf.reduce_mean( tf.boolean_mask(tf.reshape(loss_dict["object_losses"], [-1]), verb_nonpadding))) metrics["eval_verb_loss"] = tf.metrics.mean( tf.reduce_mean( tf.boolean_mask(tf.reshape(loss_dict["verbs_losses"], [-1]), verb_nonpadding)))
def _get_control_dependencies_for_print(self, labels, predictions): """Returns a list of control dependencies for tf.print. Instantiates tf.print ops, and sets the output_stream option of tf.print to print tf variables to print_dir. We print following TensorFlow variables: ["example_weights","labels","predictions","subgroups",<protected_groups>] Output is a text file at path: "<print_dir>/<tf_variable_name>.csv". The number of lines is same as number of evaluation steps. Each line is of format: "<variable_name>,<tensor_of_batch_size>". For example, "example_weight,[1.0,1.0,....]" Args: labels: A `dict` of `Tensor` Objects. Expects to have a key/value pair for the strings in self._label_column_name, self._protected_groups, and "subgroups". predictions: A `dict` of `Tensor` Objects. Expects to have a key/value pair for keys in (self._label_column_name, "class_ids"), and "example_weights". Returns: control_dependencies_for_print: A list of tf.print control_dependency ops. """ if self._print_dir: predictions_ids_print_op = tf.print( "predictions", predictions[(self._label_column_name, "class_ids")], summarize=-1, sep=",", output_stream="file:///{}/predictions.csv".format( self._print_dir)) label_print_op = tf.print( "labels", tf.reshape(labels[self._label_column_name], [-1]), summarize=-1, sep=",", output_stream="file:///{}/labels.csv".format(self._print_dir)) subgroup_print_op = tf.print( "subgroups", tf.reshape(labels["subgroup"], [-1]), summarize=-1, sep=",", output_stream="file:///{}/subgroups.csv".format( self._print_dir)) example_weights_print_op = tf.print( "example_weights", predictions[("example_weights")], summarize=-1, sep=",", output_stream="file:///{}/example_weights.csv".format( self._print_dir)) protected_col_0_print_op = tf.print( self._protected_groups[0], tf.reshape(labels[self._protected_groups[0]], [-1]), summarize=-1, sep=",", output_stream="file:///{}/{}.csv".format( self._print_dir, self._protected_groups[0])) protected_col_1_print_op = tf.print( self._protected_groups[1], tf.reshape(labels[self._protected_groups[1]], [-1]), summarize=-1, sep=",", output_stream="file:///{}/{}.csv".format( self._print_dir, self._protected_groups[1])) control_dependencies_for_print = [ predictions_ids_print_op, label_print_op, subgroup_print_op, example_weights_print_op, protected_col_0_print_op, protected_col_1_print_op ] else: control_dependencies_for_print = None return control_dependencies_for_print
def encode_decode_task(features, hparams, train, attention_weights=None): """Model core graph for the one-shot action. Args: features: a dictionary contains "inputs" that is a tensor in shape of [batch_size, num_tokens], "verb_id_seq" that is in shape of [batch_size, num_actions], "object_spans" and "param_span" tensor in shape of [batch_size, num_actions, 2]. 0 is used as padding or non-existent values. hparams: the general hyperparameters for the model. train: the train mode. attention_weights: the dict to keep attention weights for analysis. Returns: loss_dict: the losses for training. prediction_dict: the predictions for action tuples. areas: the area encodings of the task. scope: the embedding scope. """ del train input_embeddings, scope = common_embed.embed_tokens( features["task"], hparams.task_vocab_size, hparams.hidden_size, hparams) with tf.variable_scope("encode_decode", reuse=tf.AUTO_REUSE): encoder_nonpadding = tf.minimum(tf.to_float(features["task"]), 1.0) input_embeddings = tf.multiply(tf.expand_dims(encoder_nonpadding, 2), input_embeddings) encoder_input, self_attention_bias, encoder_decoder_attention_bias = ( transformer.transformer_prepare_encoder(input_embeddings, None, hparams, features=None)) encoder_input = tf.nn.dropout(encoder_input, keep_prob=1.0 - hparams.layer_prepostprocess_dropout) if hparams.instruction_encoder == "transformer": encoder_output = transformer.transformer_encoder( encoder_input, self_attention_bias, hparams, save_weights_to=attention_weights, make_image_summary=not common_layers.is_xla_compiled()) else: raise ValueError("Unsupported instruction encoder %s" % (hparams.instruction_encoder)) span_rep = hparams.get("span_rep", "area") area_encodings, area_starts, area_ends = area_utils.compute_sum_image( encoder_output, max_area_width=hparams.max_span) current_shape = tf.shape(area_encodings) if span_rep == "area": area_encodings, _, _ = area_utils.compute_sum_image( encoder_output, max_area_width=hparams.max_span) elif span_rep == "basic": area_encodings = area_utils.compute_alternative_span_rep( encoder_output, input_embeddings, max_area_width=hparams.max_span, hidden_size=hparams.hidden_size, advanced=False) elif span_rep == "coref": area_encodings = area_utils.compute_alternative_span_rep( encoder_output, input_embeddings, max_area_width=hparams.max_span, hidden_size=hparams.hidden_size, advanced=True) else: raise ValueError("xyz") areas = {} areas["encodings"] = area_encodings areas["starts"] = area_starts areas["ends"] = area_ends with tf.control_dependencies([ tf.print("encoder_output", tf.shape(encoder_output)), tf.assert_equal(current_shape, tf.shape(area_encodings), summarize=100) ]): paddings = tf.cast(tf.less(self_attention_bias, -1), tf.int32) padding_sum, _, _ = area_utils.compute_sum_image( tf.expand_dims(tf.squeeze(paddings, [1, 2]), 2), max_area_width=hparams.max_span) num_areas = common_layers.shape_list(area_encodings)[1] area_paddings = tf.reshape(tf.minimum(tf.to_float(padding_sum), 1.0), [-1, num_areas]) areas["bias"] = area_paddings decoder_nonpadding = tf.to_float( tf.greater(features["verb_refs"][:, :, 1], features["verb_refs"][:, :, 0])) if hparams.instruction_encoder == "lstm": hparams_decoder = copy.copy(hparams) hparams_decoder.set_hparam("pos", "none") else: hparams_decoder = hparams decoder_input, decoder_self_attention_bias = _prepare_decoder_input( area_encodings, decoder_nonpadding, features, hparams_decoder, embed_scope=scope) decoder_input = tf.nn.dropout(decoder_input, keep_prob=1.0 - hparams.layer_prepostprocess_dropout) if hparams.instruction_decoder == "transformer": decoder_output = transformer.transformer_decoder( decoder_input=decoder_input, encoder_output=encoder_output, decoder_self_attention_bias=decoder_self_attention_bias, encoder_decoder_attention_bias=encoder_decoder_attention_bias, hparams=hparams_decoder) else: raise ValueError("Unsupported instruction encoder %s" % (hparams.instruction_encoder)) return decoder_output, decoder_nonpadding, areas, scope
def create_optimizer(config, loss, num_train_steps): """Creates an optimizer training op.""" global_step = tf.train.get_or_create_global_step() learning_rate = tf.constant(value=config.learning_rate, shape=[], dtype=tf.float32) # Implements linear decay of the learning rate. learning_rate = tf.train.polynomial_decay(learning_rate, global_step, num_train_steps, end_learning_rate=0.0, power=1.0, cycle=False) # Implements linear warmup. I.e., if global_step < num_warmup_steps, the # learning rate will be `global_step/num_warmup_steps * init_lr`. num_warmup_steps = int(num_train_steps * config.warmup_proportion) if num_warmup_steps: global_steps_int = tf.cast(global_step, tf.int32) warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32) global_steps_float = tf.cast(global_steps_int, tf.float32) warmup_steps_float = tf.cast(warmup_steps_int, tf.float32) warmup_percent_done = global_steps_float / warmup_steps_float warmup_learning_rate = config.learning_rate * warmup_percent_done is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32) learning_rate = ((1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate) # It is recommended that you use this optimizer for fine tuning, since this # is how the model was trained (note that the Adam m/v variables are NOT # loaded from init_checkpoint.) if config.lr_decay > 0: learning_rate = _get_edecay_lrs(config, learning_rate) optimizer = AdamWeightDecayOptimizer( learning_rate=learning_rate, weight_decay_rate=config.weight_decay_rate, beta_1=0.9, beta_2=0.999, epsilon=1e-6, exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"]) #print(learning_rate) #if config.use_tpu: # optimizer = contrib_tpu.CrossShardOptimizer(optimizer) tvars = tf.trainable_variables() grads = tf.gradients(loss, tvars) #print(tvars) # This is how the model was pre-trained. (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0) train_op = optimizer.apply_gradients(zip(grads, tvars), global_step=global_step) # Normally the global step update is done inside of `apply_gradients`. # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use # a different optimizer, you should probably take this line out. new_global_step = global_step + 1 tf.print("New global step") tf.print(new_global_step) train_op = tf.group(train_op, [global_step.assign(new_global_step)]) return train_op