return loss, h optimizer = EagerOptimizer(loss, optim="adam", lr=args.lr) for epoch in range(args.epochs): loss_accum = 0. step = 0 start = time.time() h = None SET_TRAIN_FLAG(True) for x, y in train_input_fn(): # Optimize the model loss_value, h = optimizer.update_with_hidden(model, h, x, y) loss_accum += loss_value step += 1 print('training time {}'.format(time.time() - start)) mean_loss = loss_accum / step print('Training Loss {}, Perplexity {}'.format(mean_loss, np.exp(mean_loss))) step = 0 loss_accum = 0 SET_TRAIN_FLAG(False) for x, y in eval_input_fn(): # Track progress # compare predicted label to actual label
class LanguageModelTrainerDistributedTf(Trainer): """A Trainer for LM distributed eager training """ def __init__(self, model_params, **kwargs): super().__init__() if type(model_params) is dict: self.model = create_model_for('lm', **model_params) else: self.model = model_params loss_fn = loss_with_state if self.model.requires_state else loss_without_state self.optimizer = EagerOptimizer(loss_fn, **kwargs) self.nsteps = kwargs.get('nsteps', 500) self._checkpoint = tf.train.Checkpoint( optimizer=self.optimizer.optimizer, model=self.model) checkpoint_dir = '{}-{}'.format("./tf-lm", os.getpid()) self.checkpoint_manager = tf.train.CheckpointManager( self._checkpoint, directory=checkpoint_dir, max_to_keep=5) strategy_type = kwargs.get('strategy_type', 'mirror') gpus = int(kwargs.get('gpus', 1)) endpoint = kwargs.get('endpoint') self.strategy = create_distribute_strategy(strategy_type, gpus, endpoint) def checkpoint(self): """This method saves a checkpoint :return: None """ self.checkpoint_manager.save() def recover_last_checkpoint(self): """Recover the last saved checkpoint :return: None """ print( self._checkpoint.restore( self.checkpoint_manager.latest_checkpoint)) @staticmethod def _num_toks(y): return tf.reduce_prod(get_shape_as_list(y)) def train(self, ts, reporting_fns, steps=0, dataset=True): """Train by looping over the steps For a `tf.dataset`-backed `fit_func`, we are using the previously wired `dataset`s in the model (and `dataset` is `True`). For `feed_dict`, we convert the ts samples to `feed_dict`s and hand them in one-by-one :param ts: The training set :param reporting_fns: A list of reporting hooks :param dataset: (`bool`) Are we using `tf.dataset`s :return: Metrics """ strategy = self.strategy def _replicated_train_step_no_state(inputs): features, y = inputs per_replica_loss = self.optimizer.update(self.model, features, y) per_replica_toks = self._num_toks(y) per_replica_report_loss = per_replica_loss * tf.cast( per_replica_toks, tf.float32) return per_replica_report_loss, per_replica_toks def _replicated_train_step_with_state(inputs, hidden): features, y = inputs per_replica_loss, new_hidden = self.optimizer.update_with_hidden( self.model, hidden, features, y) per_replica_toks = self._num_toks(y) per_replica_report_loss = per_replica_loss * tf.cast( per_replica_toks, tf.float32) return new_hidden, per_replica_report_loss, per_replica_toks with strategy.scope(): train_iter = iter(ts) SET_TRAIN_FLAG(True) epoch_loss = tf.Variable(0.0) epoch_div = tf.Variable(0, dtype=tf.int32) nstep_loss = tf.Variable(0.0) nstep_div = tf.Variable(0, dtype=tf.int32) self.nstep_start = time.time() start = time.time() @tf.function def _distributed_train_no_state(inputs): per_replica_loss, per_replica_toks = strategy.experimental_run_v2( _replicated_train_step_no_state, args=(inputs, )) return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None), strategy.reduce( tf.distribute.ReduceOp.SUM, per_replica_toks, axis=None) @tf.function def _distributed_train_with_state(inputs, hidden): h, per_replica_loss, per_replica_toks = strategy.experimental_run_v2( _replicated_train_step_with_state, args=( inputs, hidden, )) step_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None) step_toks = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_toks, axis=None) return h, step_loss, step_toks h = None for i in range(steps): inputs = next(train_iter) if self.model.requires_state: h, step_loss, step_toks = _distributed_train_with_state( inputs, h) else: step_loss, step_toks = _distributed_train_no_state(inputs) epoch_loss.assign_add(step_loss) nstep_loss.assign_add(step_loss) epoch_div.assign_add(step_toks) nstep_div.assign_add(step_toks) step = self.optimizer.global_step.numpy() + 1 if step % self.nsteps == 0: metrics = self.calc_metrics(nstep_loss.numpy(), nstep_div.numpy()) self.report(step, metrics, self.nstep_start, 'Train', 'STEP', reporting_fns, self.nsteps) nstep_loss.assign(0.0) nstep_div.assign(0) self.nstep_start = time.time() epoch_loss = epoch_loss.numpy() epoch_div = epoch_div.numpy() metrics = self.calc_metrics(epoch_loss, epoch_div) self.train_epochs += 1 self.report(self.train_epochs, metrics, start, 'Train', 'EPOCH', reporting_fns) return metrics def calc_metrics(self, agg, norm): metrics = super().calc_metrics(agg, norm) metrics['perplexity'] = np.exp(metrics['avg_loss']) return metrics def test(self, vs, reporting_fns, phase, steps=0): """Run an epoch of testing over the dataset If we are using a `tf.dataset`-based `fit_func`, we will just cycle the number of steps and let the `dataset` yield new batches. If we are using `feed_dict`s, we convert each batch from the `DataFeed` and pass that into TF as the `feed_dict` :param vs: A validation set :param reporting_fns: Reporting hooks :param phase: The phase of evaluation (`Test`, `Valid`) :param dataset: (`bool`) Are we using `tf.dataset`s :return: Metrics """ strategy = self.strategy def _replicated_test_step_no_state(inputs): features, y = inputs per_replica_loss = loss_without_state(self.model, features, y) per_replica_toks = self._num_toks(y) per_replica_report_loss = per_replica_loss * tf.cast( per_replica_toks, tf.float32) return per_replica_report_loss, per_replica_toks def _replicated_test_step_with_state(inputs, hidden): features, y = inputs per_replica_loss, new_hidden = loss_with_state( self.model, hidden, features, y) per_replica_toks = self._num_toks(y) per_replica_report_loss = per_replica_loss * tf.cast( per_replica_toks, tf.float32) return new_hidden, per_replica_report_loss, per_replica_toks with strategy.scope(): SET_TRAIN_FLAG(False) test_iter = iter(vs) epoch_loss = tf.Variable(0.0) epoch_div = tf.Variable(0, dtype=tf.int32) self.nstep_start = time.time() start = time.time() @tf.function def _distributed_test_no_state(inputs): per_replica_loss, per_replica_toks = strategy.experimental_run_v2( _replicated_test_step_no_state, args=(inputs, )) return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None), strategy.reduce( tf.distribute.ReduceOp.SUM, per_replica_toks, axis=None) @tf.function def _distributed_test_with_state(inputs, hidden): h, per_replica_loss, per_replica_toks = strategy.experimental_run_v2( _replicated_test_step_with_state, args=( inputs, hidden, )) step_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None) step_toks = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_toks, axis=None) return h, step_loss, step_toks epochs = 0 if phase == 'Valid': self.valid_epochs += 1 epochs = self.valid_epochs h = None for i in range(steps): inputs = next(test_iter) if self.model.requires_state: h, per_replica_loss, per_replica_toks = _distributed_test_with_state( inputs, h) else: per_replica_loss, per_replica_toks = _distributed_test_no_state( inputs) epoch_loss.assign_add(per_replica_loss) epoch_div.assign_add(per_replica_toks) metrics = self.calc_metrics(epoch_loss.numpy(), epoch_div.numpy()) self.report(epochs, metrics, start, phase, 'EPOCH', reporting_fns) return metrics def distribute(self, dataset): return self.strategy.experimental_distribute_dataset(dataset)
class LanguageModelTrainerEagerTf(Trainer): """A Trainer to use for eager mode """ def __init__(self, model_params, **kwargs): super().__init__() if type(model_params) is dict: self.model = create_model_for('lm', **model_params) else: self.model = model_params loss_fn = loss_with_state if self.model.requires_state else loss_without_state self.optimizer = EagerOptimizer(loss_fn, **kwargs) self.nsteps = kwargs.get('nsteps', 500) self._checkpoint = tf.train.Checkpoint( optimizer=self.optimizer.optimizer, model=self.model) checkpoint_dir = '{}-{}'.format("./tf-lm", os.getpid()) self.checkpoint_manager = tf.train.CheckpointManager( self._checkpoint, directory=checkpoint_dir, max_to_keep=5) def checkpoint(self): """This method saves a checkpoint :return: None """ self.checkpoint_manager.save() def recover_last_checkpoint(self): """Recover the last saved checkpoint :return: None """ print( self._checkpoint.restore( self.checkpoint_manager.latest_checkpoint)) @staticmethod def _num_toks(y): return tf.reduce_prod(get_shape_as_list(y)) def train(self, ts, reporting_fns): """Train by looping over the steps For a `tf.dataset`-backed `fit_func`, we are using the previously wired `dataset`s in the model (and `dataset` is `True`). For `feed_dict`, we convert the ts samples to `feed_dict`s and hand them in one-by-one :param ts: The training set :param reporting_fns: A list of reporting hooks :param dataset: (`bool`) Are we using `tf.dataset`s :return: Metrics """ SET_TRAIN_FLAG(True) epoch_loss = tf.Variable(0.0) epoch_div = tf.Variable(0, dtype=tf.int32) nstep_loss = tf.Variable(0.0) nstep_div = tf.Variable(0, dtype=tf.int32) self.nstep_start = time.perf_counter() start = time.perf_counter() def _train_step_no_state(inputs): """Replicated training step.""" features, y = inputs loss = self.optimizer.update(self.model, features, y) toks = self._num_toks(y) report_loss = loss * tf.cast(toks, tf.float32) return report_loss, toks def _train_step_with_state(inputs, hidden): """Replicated training step.""" features, y = inputs loss, hidden = self.optimizer.update_with_hidden( self.model, hidden, features, y) toks = self._num_toks(y) report_loss = loss * tf.cast(toks, tf.float32) return hidden, report_loss, toks if get_version(tf) >= 2: _train_step_with_state = tf.function(_train_step_with_state) _train_step_no_state = tf.function(_train_step_no_state) h = None for inputs in ts: if self.model.requires_state: h, step_report_loss, step_toks = _train_step_with_state( inputs, h) else: step_report_loss, step_toks = _train_step_no_state(inputs) epoch_loss.assign_add(step_report_loss) nstep_loss.assign_add(step_report_loss) epoch_div.assign_add(step_toks) nstep_div.assign_add(step_toks) step = self.optimizer.global_step.numpy() + 1 if step % self.nsteps == 0: metrics = self.calc_metrics(nstep_loss.numpy(), nstep_div.numpy()) self.report(step, metrics, self.nstep_start, 'Train', 'STEP', reporting_fns, self.nsteps) nstep_loss.assign(0.0) nstep_div.assign(0) self.nstep_start = time.perf_counter() epoch_loss = epoch_loss.numpy() epoch_div = epoch_div.numpy() metrics = self.calc_metrics(epoch_loss, epoch_div) self.train_epochs += 1 self.report(self.train_epochs, metrics, start, 'Train', 'EPOCH', reporting_fns) return metrics def calc_metrics(self, agg, norm): metrics = super().calc_metrics(agg, norm) metrics['perplexity'] = np.exp(metrics['avg_loss']) return metrics def test(self, vs, reporting_fns, phase): """Run an epoch of testing over the dataset If we are using a `tf.dataset`-based `fit_func`, we will just cycle the number of steps and let the `dataset` yield new batches. If we are using `feed_dict`s, we convert each batch from the `DataFeed` and pass that into TF as the `feed_dict` :param vs: A validation set :param reporting_fns: Reporting hooks :param phase: The phase of evaluation (`Test`, `Valid`) :param dataset: (`bool`) Are we using `tf.dataset`s :return: Metrics """ total_loss = 0.0 total_toks = 0 epochs = 0 if phase == 'Valid': self.valid_epochs += 1 epochs = self.valid_epochs SET_TRAIN_FLAG(False) start = time.perf_counter() h = None for features, y in vs: if self.model.requires_state: loss_value, h = loss_with_state(self.model, h, features, y) else: loss_value = loss_without_state(self.model, features, y) loss_value = loss_value.numpy() toks = self._num_toks(y) total_loss += loss_value * tf.cast(toks, tf.float32).numpy() total_toks += toks.numpy() metrics = self.calc_metrics(total_loss, total_toks) self.report(epochs, metrics, start, phase, 'EPOCH', reporting_fns) return metrics