Method train

lit_nlp/examples/glue/models.py:185–229

Run fine-tuning.

(
      self,
      train_inputs: list[JsonDict],
      validation_inputs: list[JsonDict],
      learning_rate=2e-5,
      batch_size=32,
      num_epochs=3,
      keras_callbacks=None,
  )

    return tf.data.Dataset.from_tensor_slices((dict(encoded_input), labels))

  def train(
      self,
      train_inputs: list[JsonDict],
      validation_inputs: list[JsonDict],
      learning_rate=2e-5,
      batch_size=32,
      num_epochs=3,
      keras_callbacks=None,
  ):
    """Run fine-tuning."""
    train_dataset = (
        self._make_dataset(train_inputs)
        .shuffle(128)
        .batch(batch_size)
        .repeat(-1)
    )
    # Use larger batch for validation since inference is about 1/2 memory usage
    # of backprop.
    eval_batch_size = 2 * batch_size
    validation_dataset = self._make_dataset(validation_inputs).batch(
        eval_batch_size
    )

    # Prepare model for training.
    opt = keras.optimizers.Adam(learning_rate=learning_rate, epsilon=1e-08)
    if self.is_regression:
      loss = keras.losses.MeanSquaredError()
      metric = keras.metrics.RootMeanSquaredError("rmse")
    else:
      loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
      metric = keras.metrics.SparseCategoricalAccuracy("accuracy")
    self.model.compile(optimizer=opt, loss=loss, metrics=[metric])

    steps_per_epoch = len(train_inputs) // batch_size
    validation_steps = len(validation_inputs) // eval_batch_size
    history = self.model.fit(
        train_dataset,
        epochs=num_epochs,
        steps_per_epoch=steps_per_epoch,
        validation_data=validation_dataset,
        validation_steps=validation_steps,
        callbacks=keras_callbacks,
        verbose=2,
    )
    return history

  def save(self, path: str):
    """Save model weights and tokenizer info.