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