def test_create_dataset_from_tables(self):
testdata_dir = 'poem/testdata' # Assume $PWD == "google_research/".
dataset = pipelines.create_dataset_from_tables(
[os.path.join(FLAGS.test_srcdir, testdata_dir, 'tfe-2.tfrecords')],
batch_sizes=[4],
num_instances_per_record=2,
shuffle=True,
num_epochs=None,
keypoint_names_3d=keypoint_profiles.create_keypoint_profile_or_die(
'LEGACY_3DH36M17').keypoint_names,
keypoint_names_2d=keypoint_profiles.create_keypoint_profile_or_die(
'LEGACY_2DCOCO13').keypoint_names,
seed=0)
inputs = list(dataset.take(1))[0]
self.assertCountEqual(inputs.keys(), [
'image_sizes', 'keypoints_2d', 'keypoint_scores_2d', 'keypoints_3d'
])
self.assertEqual(inputs['image_sizes'].shape, [4, 2, 2])
self.assertEqual(inputs['keypoints_2d'].shape, [4, 2, 13, 2])
self.assertEqual(inputs['keypoint_scores_2d'].shape, [4, 2, 13])
self.assertEqual(inputs['keypoints_3d'].shape, [4, 2, 17, 3])
def run(input_dataset_class, common_module, keypoint_profiles_module,
input_example_parser_creator, keypoint_preprocessor_3d):
"""Runs training pipeline.
Args:
input_dataset_class: An input dataset class that matches input table type.
common_module: A Python module that defines common flags and constants.
keypoint_profiles_module: A Python module that defines keypoint profiles.
input_example_parser_creator: A function handle for creating data parser
function. If None, uses the default parser creator.
keypoint_preprocessor_3d: A function handle for preprocessing raw 3D
keypoints.
"""
_validate(common_module)
log_dir_path = FLAGS.log_dir_path
pipeline_utils.create_dir_and_save_flags(flags, log_dir_path,
'all_flags.train.json')
# Setup summary writer.
summary_writer = tf.summary.create_file_writer(
os.path.join(log_dir_path, 'train_logs'), flush_millis=10000)
# Setup configuration.
keypoint_profile_2d = keypoint_profiles_module.create_keypoint_profile_or_die(
FLAGS.keypoint_profile_name_2d)
keypoint_profile_3d = keypoint_profiles_module.create_keypoint_profile_or_die(
FLAGS.keypoint_profile_name_3d)
model = algorithms.get_algorithm(
algorithm_type=FLAGS.algorithm_type,
pose_embedding_dim=FLAGS.pose_embedding_dim,
view_embedding_dim=FLAGS.view_embedding_dim,
fusion_op_type=FLAGS.fusion_op_type,
view_loss_weight=FLAGS.view_loss_weight,
regularization_loss_weight=FLAGS.regularization_loss_weight,
disentangle_loss_weight=FLAGS.disentangle_loss_weight,
embedder_type=FLAGS.embedder_type)
optimizers = algorithms.get_optimizers(
algorithm_type=FLAGS.algorithm_type, learning_rate=FLAGS.learning_rate)
global_step = optimizers['encoder_optimizer'].iterations
ckpt_manager, _, _ = utils.create_checkpoint(
log_dir_path, **optimizers, model=model, global_step=global_step)
# Setup the training dataset.
dataset = pipelines.create_dataset_from_tables(
[FLAGS.input_table], [FLAGS.batch_size],
num_instances_per_record=2,
shuffle=True,
num_epochs=None,
drop_remainder=True,
keypoint_names_2d=keypoint_profile_2d.keypoint_names,
keypoint_names_3d=keypoint_profile_3d.keypoint_names,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
dataset_class=input_dataset_class,
input_example_parser_creator=input_example_parser_creator)
def train_one_iteration(inputs):
"""Trains the model for one iteration.
Args:
inputs: A dictionary for training inputs.
Returns:
The training loss for this iteration.
"""
_, side_outputs = pipelines.create_model_input(
inputs, FLAGS.model_input_keypoint_type, keypoint_profile_2d,
keypoint_profile_3d)
keypoints_2d = side_outputs[common_module.KEY_PREPROCESSED_KEYPOINTS_2D]
keypoints_3d, _ = keypoint_preprocessor_3d(
inputs[common_module.KEY_KEYPOINTS_3D],
keypoint_profile_3d,
normalize_keypoints_3d=True)
keypoints_2d, keypoints_3d = data_utils.shuffle_batches(
[keypoints_2d, keypoints_3d])
return model.train((keypoints_2d, keypoints_3d), **optimizers)
if FLAGS.compile:
train_one_iteration = tf.function(train_one_iteration)
record_every_n_steps = min(100, FLAGS.num_iterations)
save_ckpt_every_n_steps = min(10000, FLAGS.num_iterations)
with summary_writer.as_default():
with tf.summary.record_if(global_step % record_every_n_steps == 0):
start = time.time()
for inputs in dataset:
if global_step >= FLAGS.num_iterations:
break
model_losses = train_one_iteration(inputs)
duration = time.time() - start
start = time.time()
for tag, losses in model_losses.items():
for name, loss in losses.items():
tf.summary.scalar(
'train/{}/{}'.format(tag, name), loss, step=global_step)
for tag, optimizer in optimizers.items():
tf.summary.scalar(
'train/{}_learning_rate'.format(tag),
optimizer.lr,
step=global_step)
tf.summary.scalar('train/batch_time', duration, step=global_step)
tf.summary.scalar('global_step/sec', 1 / duration, step=global_step)
if global_step % record_every_n_steps == 0:
logging.info('Iter[{}/{}], {:.6f}s/iter, loss: {:.4f}'.format(
global_step.numpy(), FLAGS.num_iterations, duration,
model_losses['encoder']['total_loss'].numpy()))
# Save checkpoint.
if global_step % save_ckpt_every_n_steps == 0:
ckpt_manager.save(checkpoint_number=global_step)
logging.info('Checkpoint saved at step %d.', global_step.numpy())
def run(input_dataset_class, common_module, keypoint_profiles_module,
input_example_parser_creator):
"""Runs training pipeline.
Args:
input_dataset_class: An input dataset class that matches input table type.
common_module: A Python module that defines common flags and constants.
keypoint_profiles_module: A Python module that defines keypoint profiles.
input_example_parser_creator: A function handle for creating data parser
function. If None, uses the default parser creator.
"""
log_dir_path = FLAGS.log_dir_path
if not tf.io.gfile.exists(log_dir_path):
raise ValueError(
'The directory {} does not exist. Please provide a new log_dir_path.'
.format(log_dir_path))
eval_log_dir = os.path.join(log_dir_path, FLAGS.eval_name)
pipeline_utils.create_dir_and_save_flags(flags, eval_log_dir,
'all_flags.eval_with_features.json')
# Setup summary writer.
summary_writer = tf.summary.create_file_writer(
eval_log_dir, flush_millis=10000)
# Setup configuration.
keypoint_profile_2d = keypoint_profiles_module.create_keypoint_profile_or_die(
FLAGS.keypoint_profile_name_2d)
# Setup model.
input_length = math.ceil(FLAGS.num_frames / FLAGS.downsample_rate)
if FLAGS.input_features_dim > 0:
feature_dim = FLAGS.input_features_dim
input_shape = (input_length, feature_dim)
else:
feature_dim = None
input_shape = (input_length, 13 * 2)
classifier = models.get_temporal_classifier(
FLAGS.classifier_type,
input_shape=input_shape,
num_classes=FLAGS.num_classes)
ema_classifier = tf.keras.models.clone_model(classifier)
global_step = tf.Variable(0, dtype=tf.int64, trainable=False)
dataset = pipelines.create_dataset_from_tables(
FLAGS.input_tables, [int(x) for x in FLAGS.batch_sizes],
num_instances_per_record=1,
shuffle=False,
num_epochs=1,
drop_remainder=False,
keypoint_names_2d=keypoint_profile_2d.keypoint_names,
feature_dim=feature_dim,
num_classes=FLAGS.num_classes,
num_frames=FLAGS.num_frames,
common_module=common_module,
dataset_class=input_dataset_class,
input_example_parser_creator=input_example_parser_creator)
if FLAGS.compile:
classifier.call = tf.function(classifier.call)
ema_classifier.call = tf.function(ema_classifier.call)
top_1_best_accuracy = None
top_5_best_accuracy = None
evaluated_last_ckpt = False
def timeout_fn():
"""Timeout function to stop the evaluation."""
return evaluated_last_ckpt
def evaluate_once():
"""Evaluates the model for one time."""
_, status, _ = utils.create_checkpoint(
log_dir_path,
model=classifier,
ema_model=ema_classifier,
global_step=global_step)
status.expect_partial()
logging.info('Last checkpoint [iteration: %d] restored at %s.',
global_step.numpy(), log_dir_path)
if global_step.numpy() >= FLAGS.max_iteration:
nonlocal evaluated_last_ckpt
evaluated_last_ckpt = True
top_1_accuracy = tf.keras.metrics.CategoricalAccuracy()
top_5_accuracy = tf.keras.metrics.TopKCategoricalAccuracy(k=5)
for inputs in dataset:
if FLAGS.input_features_dim > 0:
features = inputs[common_module.KEY_FEATURES]
else:
features, _ = pipelines.create_model_input(
inputs, common_module.MODEL_INPUT_KEYPOINT_TYPE_2D_INPUT,
keypoint_profile_2d)
features = tf.squeeze(features, axis=1)
features = features[:, ::FLAGS.downsample_rate, :]
labels = inputs[common_module.KEY_CLASS_TARGETS]
labels = tf.squeeze(labels, axis=1)
if FLAGS.use_moving_average:
outputs = ema_classifier(features, training=False)
else:
outputs = classifier(features, training=False)
top_1_accuracy.update_state(y_true=labels, y_pred=outputs)
top_5_accuracy.update_state(y_true=labels, y_pred=outputs)
nonlocal top_1_best_accuracy
if (top_1_best_accuracy is None or
top_1_accuracy.result().numpy() > top_1_best_accuracy):
top_1_best_accuracy = top_1_accuracy.result().numpy()
nonlocal top_5_best_accuracy
if (top_5_best_accuracy is None or
top_5_accuracy.result().numpy() > top_5_best_accuracy):
top_5_best_accuracy = top_5_accuracy.result().numpy()
tf.summary.scalar(
'eval/Basic/Top1_Accuracy',
top_1_accuracy.result(),
step=global_step.numpy())
tf.summary.scalar(
'eval/Best/Top1_Accuracy',
top_1_best_accuracy,
step=global_step.numpy())
tf.summary.scalar(
'eval/Basic/Top5_Accuracy',
top_5_accuracy.result(),
step=global_step.numpy())
tf.summary.scalar(
'eval/Best/Top5_Accuracy',
top_5_best_accuracy,
step=global_step.numpy())
logging.info('Accuracy: {:.2f}'.format(top_1_accuracy.result().numpy()))
with summary_writer.as_default():
with tf.summary.record_if(True):
if FLAGS.continuous_eval:
for _ in tf.train.checkpoints_iterator(log_dir_path, timeout=1,
timeout_fn=timeout_fn):
evaluate_once()
else:
evaluate_once()
def run(input_dataset_class, common_module, keypoint_profiles_module,
input_example_parser_creator):
"""Runs training pipeline.
Args:
input_dataset_class: An input dataset class that matches input table type.
common_module: A Python module that defines common flags and constants.
keypoint_profiles_module: A Python module that defines keypoint profiles.
input_example_parser_creator: A function handle for creating data parser
function. If None, uses the default parser creator.
"""
log_dir_path = FLAGS.log_dir_path
pipeline_utils.create_dir_and_save_flags(
flags, log_dir_path, 'all_flags.train_with_features.json')
# Setup summary writer.
summary_writer = tf.summary.create_file_writer(
os.path.join(log_dir_path, 'train_logs'), flush_millis=10000)
# Setup configuration.
keypoint_profile_2d = keypoint_profiles_module.create_keypoint_profile_or_die(
FLAGS.keypoint_profile_name_2d)
# Setup model.
input_length = math.ceil(FLAGS.num_frames / FLAGS.downsample_rate)
if FLAGS.input_features_dim > 0:
feature_dim = FLAGS.input_features_dim
input_shape = (input_length, feature_dim)
else:
feature_dim = None
input_shape = (input_length, 13 * 2)
classifier = models.get_temporal_classifier(
FLAGS.classifier_type,
input_shape=input_shape,
num_classes=FLAGS.num_classes)
ema_classifier = tf.keras.models.clone_model(classifier)
optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate)
optimizer = tfa_optimizers.MovingAverage(optimizer)
global_step = optimizer.iterations
ckpt_manager, _, _ = utils.create_checkpoint(
log_dir_path,
optimizer=optimizer,
model=classifier,
ema_model=ema_classifier,
global_step=global_step)
# Setup the training dataset.
dataset = pipelines.create_dataset_from_tables(
FLAGS.input_tables, [int(x) for x in FLAGS.batch_sizes],
num_instances_per_record=1,
shuffle=True,
drop_remainder=True,
num_epochs=None,
keypoint_names_2d=keypoint_profile_2d.keypoint_names,
feature_dim=feature_dim,
num_classes=FLAGS.num_classes,
num_frames=FLAGS.num_frames,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
common_module=common_module,
dataset_class=input_dataset_class,
input_example_parser_creator=input_example_parser_creator)
loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
def train_one_iteration(inputs):
"""Trains the model for one iteration.
Args:
inputs: A dictionary for training inputs.
Returns:
loss: The training loss for this iteration.
"""
if FLAGS.input_features_dim > 0:
features = inputs[common_module.KEY_FEATURES]
else:
features, _ = pipelines.create_model_input(
inputs, common_module.MODEL_INPUT_KEYPOINT_TYPE_2D_INPUT,
keypoint_profile_2d)
features = tf.squeeze(features, axis=1)
features = features[:, ::FLAGS.downsample_rate, :]
labels = inputs[common_module.KEY_CLASS_TARGETS]
labels = tf.squeeze(labels, axis=1)
with tf.GradientTape() as tape:
outputs = classifier(features, training=True)
regularization_loss = sum(classifier.losses)
crossentropy_loss = loss_object(labels, outputs)
total_loss = crossentropy_loss + regularization_loss
trainable_variables = classifier.trainable_variables
grads = tape.gradient(total_loss, trainable_variables)
optimizer.apply_gradients(zip(grads, trainable_variables))
for grad, trainable_variable in zip(grads, trainable_variables):
tf.summary.scalar(
'summarize_grads/' + trainable_variable.name,
tf.linalg.norm(grad),
step=global_step)
return dict(
total_loss=total_loss,
crossentropy_loss=crossentropy_loss,
regularization_loss=regularization_loss)
if FLAGS.compile:
train_one_iteration = tf.function(train_one_iteration)
record_every_n_steps = min(5, FLAGS.num_iterations)
save_ckpt_every_n_steps = min(500, FLAGS.num_iterations)
with summary_writer.as_default():
with tf.summary.record_if(global_step % record_every_n_steps == 0):
start = time.time()
for inputs in dataset:
if global_step >= FLAGS.num_iterations:
break
model_losses = train_one_iteration(inputs)
duration = time.time() - start
start = time.time()
for name, loss in model_losses.items():
tf.summary.scalar('train/' + name, loss, step=global_step)
tf.summary.scalar('train/learning_rate', optimizer.lr, step=global_step)
tf.summary.scalar('train/batch_time', duration, step=global_step)
tf.summary.scalar('global_step/sec', 1 / duration, step=global_step)
if global_step % record_every_n_steps == 0:
logging.info('Iter[{}/{}], {:.6f}s/iter, loss: {:.4f}'.format(
global_step.numpy(), FLAGS.num_iterations, duration,
model_losses['total_loss'].numpy()))
# Save checkpoint.
if global_step % save_ckpt_every_n_steps == 0:
utils.assign_moving_average_vars(classifier, ema_classifier,
optimizer)
ckpt_manager.save(checkpoint_number=global_step)
logging.info('Checkpoint saved at step %d.', global_step.numpy())