# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Trains on embeddings using Keras."""

import os
from absl import app
from absl import flags

import tensorflow.compat.v1 as tf

from non_semantic_speech_benchmark.eval_embedding.keras import get_data
from non_semantic_speech_benchmark.eval_embedding.keras import models

FLAGS = flags.FLAGS

flags.DEFINE_string('file_pattern', None, 'Dataset location.')
flags.DEFINE_string('en', None, 'Embedding name.')
flags.DEFINE_string('ed', None, 'Embedding dimension.')
flags.DEFINE_string('label_name', None, 'Name of label to use.')
flags.DEFINE_list('label_list', None, 'List of possible label values.')

flags.DEFINE_integer('tbs', 1, 'Hyperparameter: batch size.')
flags.DEFINE_integer('shuffle_buffer_size', None, 'shuffle_buffer_size')

flags.DEFINE_integer('nc', None, 'num_clusters')
flags.DEFINE_boolean('ubn', None, 'Whether to use batch normalization.')
flags.DEFINE_float('lr', 0.001, 'Hyperparameter: learning rate.')

flags.DEFINE_string('logdir', None,
                    'Path to directory where to store summaries.')

flags.DEFINE_integer('training_steps', 1000,
                     'The number of steps to run training for.')
flags.DEFINE_integer('measurement_store_interval', 10,
                     'The number of steps between storing objective value in '
                     'measurements.')


def train_and_report(debug=False):
  """Trains the classifier."""
  tf.logging.info('embedding_name: %s', FLAGS.en)
  tf.logging.info('Logdir: %s', FLAGS.logdir)
  tf.logging.info('Batch size: %s', FLAGS.tbs)

  with tf.Graph().as_default():
    reader = tf.data.TFRecordDataset
    ds = get_data.get_data(
        file_pattern=FLAGS.file_pattern,
        reader=reader,
        embedding_name=FLAGS.en,
        embedding_dim=FLAGS.ed,
        preaveraged=False,
        label_name=FLAGS.label_name,
        label_list=FLAGS.label_list,
        batch_size=FLAGS.tbs,
        loop_forever=True,
        shuffle=True,
        shuffle_buffer_size=FLAGS.shuffle_buffer_size)
    emb, y_onehot = ds.make_one_shot_iterator().get_next()
    emb.shape.assert_has_rank(3)
    assert emb.shape[2] == FLAGS.ed
    y_onehot.shape.assert_has_rank(2)
    assert y_onehot.shape[1] == len(FLAGS.label_list)

    # Loss and train.
    loss, train_op = make_graph(emb, y_onehot, FLAGS.ubn, FLAGS.nc)

    # Run training.
    train_ops_to_run = {
        'train_loss': loss,
        'train_op': train_op,
        'tensorboard_summary': tf.summary.merge_all()
    }
    if debug: return
    train_loop(train_ops_to_run)
    tf.logging.info('Finished training.')


def train_loop(train_ops_to_run):
  """Simple train loop."""
  summary_writer = None
  assert FLAGS.logdir
  summary_writer = tf.summary.FileWriter(FLAGS.logdir)
  saver = tf.train.Saver(max_to_keep=None)
  save_fn = os.path.join(FLAGS.logdir, 'ckpt')

  # Gstep must be in TF so that reloads work properly.
  assert 'gstep' not in train_ops_to_run
  train_ops_to_run['gstep'] = tf.train.get_or_create_global_step()
  gstep = 0

  with tf.train.SingularMonitoredSession(
      hooks=[tf.train.StepCounterHook(summary_writer=summary_writer)],
      ) as sess:
    while gstep < FLAGS.training_steps:
      train_values = sess.run(train_ops_to_run)
      # Report objective value periodically.
      # Step must be in TF so it deals with job restarts correctly.
      gstep = train_values['gstep']
      if gstep % 10 == 0:
        tf.logging.info('{{\'step\': {}, \'train_loss\': {}}}'.format(
            gstep, train_values['train_loss']))
      if gstep % FLAGS.measurement_store_interval == 0:
        tensorboard_summary = train_values['tensorboard_summary']
        summary_writer.add_summary(tensorboard_summary, gstep)
        saver.save(sess.raw_session(), save_fn, global_step=gstep)


def make_graph(emb, y_onehot, ubn=None, nc=None):
  """Make a graph on data."""
  num_classes = y_onehot.shape[1]
  model = models.get_keras_model(num_classes, ubn, num_clusters=nc)
  logits = model(emb, training=True)
  logits.shape.assert_is_compatible_with(y_onehot.shape)
  for u_op in model.updates:
    tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, u_op)

  # Loss.
  loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)(
      y_true=y_onehot,
      y_pred=logits)
  tf.summary.scalar('xent_loss', loss)

  # Gradient.
  opt = tf.train.AdamOptimizer(
      learning_rate=FLAGS.lr, beta1=0.9, beta2=0.999, epsilon=1e-8)
  update_ops = set(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
  with tf.control_dependencies(update_ops):
    total_loss = tf.identity(loss)
  var_list = tf.trainable_variables()
  assert var_list
  train_op = opt.minimize(total_loss, tf.train.get_or_create_global_step(),
                          var_list)

  return total_loss, train_op


def main(unused_argv):
  assert FLAGS.file_pattern
  assert FLAGS.shuffle_buffer_size
  assert FLAGS.en
  assert FLAGS.ed
  assert FLAGS.label_name
  assert FLAGS.label_list
  assert FLAGS.logdir
  train_and_report()


if __name__ == '__main__':
  app.run(main)

# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Trains on embeddings using Keras."""

import os
from absl import app
from absl import flags

import tensorflow.compat.v1 as tf

from non_semantic_speech_benchmark.eval_embedding.keras import get_data
from non_semantic_speech_benchmark.eval_embedding.keras import models

FLAGS = flags.FLAGS

flags.DEFINE_string('file_pattern', None, 'Dataset location.')
flags.DEFINE_string('en', None, 'Embedding name.')
flags.DEFINE_string('ed', None, 'Embedding dimension.')
flags.DEFINE_string('label_name', None, 'Name of label to use.')
flags.DEFINE_list('label_list', None, 'List of possible label values.')

flags.DEFINE_integer('tbs', 1, 'Hyperparameter: batch size.')
flags.DEFINE_integer('shuffle_buffer_size', None, 'shuffle_buffer_size')

flags.DEFINE_integer('nc', None, 'num_clusters')
flags.DEFINE_boolean('ubn', None, 'Whether to use batch normalization.')
flags.DEFINE_float('lr', 0.001, 'Hyperparameter: learning rate.')

flags.DEFINE_string('logdir', None,
                    'Path to directory where to store summaries.')

flags.DEFINE_integer('training_steps', 1000,
                     'The number of steps to run training for.')
flags.DEFINE_integer('measurement_store_interval', 10,
                     'The number of steps between storing objective value in '
                     'measurements.')


def train_and_report(debug=False):
  """Trains the classifier."""
  tf.logging.info('embedding_name: %s', FLAGS.en)
  tf.logging.info('Logdir: %s', FLAGS.logdir)
  tf.logging.info('Batch size: %s', FLAGS.tbs)

  with tf.Graph().as_default():
    reader = tf.data.TFRecordDataset
    ds = get_data.get_data(
        file_pattern=FLAGS.file_pattern,
        reader=reader,
        embedding_name=FLAGS.en,
        embedding_dim=FLAGS.ed,
        preaveraged=False,
        label_name=FLAGS.label_name,
        label_list=FLAGS.label_list,
        batch_size=FLAGS.tbs,
        loop_forever=True,
        shuffle=True,
        shuffle_buffer_size=FLAGS.shuffle_buffer_size)
    emb, y_onehot = ds.make_one_shot_iterator().get_next()
    emb.shape.assert_has_rank(3)
    assert emb.shape[2] == FLAGS.ed
    y_onehot.shape.assert_has_rank(2)
    assert y_onehot.shape[1] == len(FLAGS.label_list)

    # Loss and train.
    loss, train_op = make_graph(emb, y_onehot, FLAGS.ubn, FLAGS.nc)

    # Run training.
    train_ops_to_run = {
        'train_loss': loss,
        'train_op': train_op,
        'tensorboard_summary': tf.summary.merge_all()
    }
    if debug: return
    train_loop(train_ops_to_run)
    tf.logging.info('Finished training.')


def train_loop(train_ops_to_run):
  """Simple train loop."""
  summary_writer = None
  assert FLAGS.logdir
  summary_writer = tf.summary.FileWriter(FLAGS.logdir)
  saver = tf.train.Saver(max_to_keep=None)
  save_fn = os.path.join(FLAGS.logdir, 'ckpt')

  # Gstep must be in TF so that reloads work properly.
  assert 'gstep' not in train_ops_to_run
  train_ops_to_run['gstep'] = tf.train.get_or_create_global_step()
  gstep = 0

  with tf.train.SingularMonitoredSession(
      hooks=[tf.train.StepCounterHook(summary_writer=summary_writer)],
      ) as sess:
    while gstep < FLAGS.training_steps:
      train_values = sess.run(train_ops_to_run)
      # Report objective value periodically.
      # Step must be in TF so it deals with job restarts correctly.
      gstep = train_values['gstep']
      if gstep % 10 == 0:
        tf.logging.info('{{\'step\': {}, \'train_loss\': {}}}'.format(
            gstep, train_values['train_loss']))
      if gstep % FLAGS.measurement_store_interval == 0:
        tensorboard_summary = train_values['tensorboard_summary']
        summary_writer.add_summary(tensorboard_summary, gstep)
        saver.save(sess.raw_session(), save_fn, global_step=gstep)


def make_graph(emb, y_onehot, ubn=None, nc=None):
  """Make a graph on data."""
  num_classes = y_onehot.shape[1]
  model = models.get_keras_model(num_classes, ubn, num_clusters=nc)
  logits = model(emb, training=True)
  logits.shape.assert_is_compatible_with(y_onehot.shape)
  for u_op in model.updates:
    tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, u_op)

  # Loss.
  loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)(
      y_true=y_onehot,
      y_pred=logits)
  tf.summary.scalar('xent_loss', loss)

  # Gradient.
  opt = tf.train.AdamOptimizer(
      learning_rate=FLAGS.lr, beta1=0.9, beta2=0.999, epsilon=1e-8)
  update_ops = set(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
  with tf.control_dependencies(update_ops):
    total_loss = tf.identity(loss)
  var_list = tf.trainable_variables()
  assert var_list
  train_op = opt.minimize(total_loss, tf.train.get_or_create_global_step(),
                          var_list)

  return total_loss, train_op


def main(unused_argv):
  assert FLAGS.file_pattern
  assert FLAGS.shuffle_buffer_size
  assert FLAGS.en
  assert FLAGS.ed
  assert FLAGS.label_name
  assert FLAGS.label_list
  assert FLAGS.logdir
  train_and_report()


if __name__ == '__main__':
  app.run(main)

# coding=utf-8 # Copyright 2020 The Google Research Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Lint as: python3 Trains on embeddings using Keras. Trains the classifier. # Loss and train. # Run training. Simple train loop. # Gstep must be in TF so that reloads work properly. # Report objective value periodically. # Step must be in TF so it deals with job restarts correctly. Make a graph on data. # Loss. # Gradient.