training_utils.py 11.4 KB
Newer Older
Frederick Liu's avatar
Frederick Liu committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Hongkun Yu's avatar
Hongkun Yu committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Frederick Liu's avatar
Frederick Liu committed
14

Hongkun Yu's avatar
Hongkun Yu committed
15
"""XLNet training utils."""
Hongkun Yu's avatar
Hongkun Yu committed
16
17
18

import os
import re
Frederick Liu's avatar
Frederick Liu committed
19
from typing import Any, Callable, Dict, Optional, Text
Hongkun Yu's avatar
Hongkun Yu committed
20
21
22

from absl import logging
import tensorflow as tf
23

24
from official.nlp.bert import model_training_utils
Hongkun Yu's avatar
Hongkun Yu committed
25
from official.nlp.xlnet import data_utils
Frederick Liu's avatar
Frederick Liu committed
26
27
28

# pytype: disable=attribute-error
# pylint: disable=g-bare-generic,unused-import
Hongkun Yu's avatar
Hongkun Yu committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

_MIN_SUMMARY_STEPS = 10


def _save_checkpoint(checkpoint, model_dir, checkpoint_prefix):
  """Saves model to with provided checkpoint prefix."""

  checkpoint_path = os.path.join(model_dir, checkpoint_prefix)
  saved_path = checkpoint.save(checkpoint_path)
  logging.info("Saving model as TF checkpoint: %s", saved_path)
  return


def _float_metric_value(metric):
  """Gets the value of a float-value keras metric."""
  return metric.result().numpy().astype(float)


def train(
    strategy: tf.distribute.Strategy,
    model_fn: Callable,
    input_meta_data: Dict,
    train_input_fn: Callable,
    total_training_steps: int,
    steps_per_loop: int,
    optimizer: tf.keras.optimizers.Optimizer,
    learning_rate_fn: tf.keras.optimizers.schedules.LearningRateSchedule,
    eval_fn: Optional[Callable[[tf.keras.Model, int, tf.summary.SummaryWriter],
                               Any]] = None,
    metric_fn: Optional[Callable[[], tf.keras.metrics.Metric]] = None,
    init_checkpoint: Optional[Text] = None,
60
    init_from_transformerxl: Optional[bool] = False,
Hongkun Yu's avatar
Hongkun Yu committed
61
    model_dir: Optional[Text] = None,
Hongkun Yu's avatar
Hongkun Yu committed
62
63
    save_steps: Optional[int] = None,
    run_eagerly: Optional[bool] = False):
Hongkun Yu's avatar
Hongkun Yu committed
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
  """Runs customized training.

  Args:
      strategy: Distribution strategy on which to run low level training loop.
      model_fn: The function returns a keras.Model.
      input_meta_data: A dictionary of params: `mem_len`, `lr_layer_decay_rate`,
        `n_layer`, `batch_size_per_core` and `d_model`.
      train_input_fn: Function returns a tf.data.Dataset used for training.
      total_training_steps: Number of steps to train in total.
      steps_per_loop: Number of steps per graph-mode loop. In order to reduce
        communication in eager context, training logs are printed every
        steps_per_loop.
      optimizer: The optimizer for model.
      learning_rate_fn: the learning rate schedule.
      eval_fn: A callback of evaluation function, that takes a keras.Model,
        current step and evaluation summary writer.
      metric_fn: A metrics function returns a Keras Metric object to record
        evaluation result using evaluation dataset or with training dataset
        after every epoch.
      init_checkpoint: Optional checkpoint to load to `sub_model` returned by
        `model_fn`.
85
86
      init_from_transformerxl: Whether to load to `transformerxl_model` of
        `model_fn`.
Hongkun Yu's avatar
Hongkun Yu committed
87
88
      model_dir: The directory of model (checkpoints, summaries).
      save_steps: The frequency to save checkpoints. Every save_steps, we save a
89
90
        model checkpoint. Model checkpoint will be saved and evaluation will be
        conducted if evaluation dataset is provided.
Hongkun Yu's avatar
Hongkun Yu committed
91
      run_eagerly: Whether to run training eagerly.
Hongkun Yu's avatar
Hongkun Yu committed
92
93
94
95
96
97
98

  Returns:
      Last training step logits if training happens, otherwise returns None.
  Raises:
    TypeError: if model directory is not specified.
  """
  required_arguments = [
99
100
      train_input_fn, total_training_steps, steps_per_loop, optimizer,
      learning_rate_fn, save_steps
Hongkun Yu's avatar
Hongkun Yu committed
101
102
  ]
  if [arg for arg in required_arguments if arg is None]:
Hongkun Yu's avatar
Hongkun Yu committed
103
    raise ValueError("`train_input_fn`, `total_training_steps`, "
104
                     "`steps_per_loop`, `optimizer`, `save_steps` and "
Hongkun Yu's avatar
Hongkun Yu committed
105
                     "`learning_rate_fn` are required parameters.")
Hongkun Yu's avatar
Hongkun Yu committed
106
107
  if not model_dir:
    raise TypeError("Model directory must be specified.")
Hongkun Yu's avatar
Hongkun Yu committed
108
  train_iterator = data_utils.get_input_iterator(train_input_fn, strategy)
Hongkun Yu's avatar
Hongkun Yu committed
109
110
  if not tf.io.gfile.exists(model_dir):
    tf.io.gfile.mkdir(model_dir)
Hongkun Yu's avatar
Hongkun Yu committed
111
112
113
114
115
116
  # Create summary writers
  summary_dir = os.path.join(model_dir, "summaries")
  if not tf.io.gfile.exists(summary_dir):
    tf.io.gfile.mkdir(summary_dir)
  train_summary_writer = None
  eval_summary_writer = None
Hongkun Yu's avatar
Hongkun Yu committed
117
  if eval_fn:
Hongkun Yu's avatar
Hongkun Yu committed
118
    eval_summary_writer = tf.summary.create_file_writer(
Hongkun Yu's avatar
Hongkun Yu committed
119
        os.path.join(summary_dir, "eval"))
Hongkun Yu's avatar
Hongkun Yu committed
120
121
122
123
  if steps_per_loop >= _MIN_SUMMARY_STEPS:
    # Only writes summary when the stats are collected sufficiently over
    # enough steps.
    train_summary_writer = tf.summary.create_file_writer(
Hongkun Yu's avatar
Hongkun Yu committed
124
        os.path.join(summary_dir, "train"))
Hongkun Yu's avatar
Hongkun Yu committed
125
126
127
128
129
130

  with strategy.scope():
    model = model_fn()

    if init_checkpoint:
      logging.info("restore from %s", init_checkpoint)
131
132
133
134
135
      if init_from_transformerxl:
        checkpoint = tf.train.Checkpoint(
            transformer_xl=model.transformerxl_model)
      else:
        checkpoint = tf.train.Checkpoint(model=model)
Hongkun Yu's avatar
Hongkun Yu committed
136
      checkpoint.restore(init_checkpoint)
Hongkun Yu's avatar
Hongkun Yu committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

    model.optimizer = optimizer

    if not hasattr(model, "optimizer"):
      raise ValueError("User should set optimizer attribute to model.")

    train_loss_metric = tf.keras.metrics.Mean("training_loss", dtype=tf.float32)
    train_metric = None
    if metric_fn:
      train_metric = metric_fn()

    def _replicated_step(inputs, mem=None):
      """Replicated training step."""

      inputs["mems"] = mem
      with tf.GradientTape() as tape:
        mem, logits = model(inputs, training=True)
        loss = model.losses
        train_loss_metric.update_state(loss)
        if train_metric:
          train_metric.update_state(inputs["label_ids"], logits)
        scaled_loss = loss[0] * 1.0 / float(strategy.num_replicas_in_sync)

      # Collects training variables.
      tvars = model.trainable_variables
      grads = tape.gradient(scaled_loss, tvars)
      clipped, _ = tf.clip_by_global_norm(grads, clip_norm=1.0)

      if input_meta_data["lr_layer_decay_rate"] != 1.0:
        n_layer = 0
        for i in range(len(clipped)):
          m = re.search(r"model/transformer/layer_(\d+?)/", tvars[i].name)
          if not m:
            continue
          n_layer = max(n_layer, int(m.group(1)) + 1)

        for i in range(len(clipped)):
          for l in range(n_layer):
            if "model/transformer/layer_{}/".format(l) in tvars[i].name:
              abs_rate = input_meta_data["lr_layer_decay_rate"]**(
                  n_layer - 1 - l)
              clipped[i] *= abs_rate
              logging.info("Apply mult {:.4f} to layer-{} grad of {}".format(
                  abs_rate, l, tvars[i].name))
              break

      optimizer.apply_gradients(zip(clipped, tvars))
      if input_meta_data["mem_len"] > 0:
Hongkun Yu's avatar
Hongkun Yu committed
185
        return mem
Hongkun Yu's avatar
Hongkun Yu committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211

    def train_steps(iterator, steps):
      """Performs distributed training steps in a loop.

      Args:
        iterator: the distributed iterator of training datasets.
        steps: an tf.int32 integer tensor to specify number of steps to run
          inside host training loop.

      Raises:
        ValueError: Any of the arguments or tensor shapes are invalid.

      Returns:
        logits: logits computed.
      """
      if not isinstance(steps, tf.Tensor):
        raise ValueError("steps should be an Tensor. Python object may cause "
                         "retracing.")

      def cache_fn():
        """Initializes memory tensor used in XLNet pretraining."""
        mems = []
        if input_meta_data["mem_len"] > 0:
          for _ in range(input_meta_data["n_layer"]):
            zeros = tf.zeros([
                input_meta_data["batch_size_per_core"],
Allen Wang's avatar
Allen Wang committed
212
                input_meta_data["mem_len"],
Hongkun Yu's avatar
Hongkun Yu committed
213
214
215
216
217
218
219
                input_meta_data["d_model"]
            ],
                             dtype=tf.float32)
            mems.append(zeros)
        return mems

      if input_meta_data["mem_len"] > 0:
Ken Franko's avatar
Ken Franko committed
220
        mem = strategy.run(cache_fn)
Hongkun Yu's avatar
Hongkun Yu committed
221
        for _ in tf.range(steps):
Ken Franko's avatar
Ken Franko committed
222
          mem = strategy.run(
Hongkun Yu's avatar
Hongkun Yu committed
223
224
225
226
227
228
              _replicated_step, args=(
                  next(iterator),
                  mem,
              ))
      else:
        for _ in tf.range(steps):
Ken Franko's avatar
Ken Franko committed
229
          strategy.run(_replicated_step, args=(next(iterator),))
Hongkun Yu's avatar
Hongkun Yu committed
230
231
232

    if not run_eagerly:
      train_steps = tf.function(train_steps)
Hongkun Yu's avatar
Hongkun Yu committed
233
234
235
236
237
238
239
240
241
242
243
244

    logging.info("Start training...")
    checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
    latest_checkpoint_file = tf.train.latest_checkpoint(model_dir)
    if latest_checkpoint_file:
      logging.info("Checkpoint file %s found and restoring from checkpoint",
                   latest_checkpoint_file)
      checkpoint.restore(latest_checkpoint_file)
      logging.info("Loading from checkpoint file completed")

    current_step = optimizer.iterations.numpy()
    checkpoint_name = "xlnet_step_{step}.ckpt"
Hongkun Yu's avatar
Hongkun Yu committed
245

Hongkun Yu's avatar
Hongkun Yu committed
246
247
248
249
250
    while current_step < total_training_steps:
      train_loss_metric.reset_states()
      if train_metric:
        train_metric.reset_states()

251
252
      steps = model_training_utils.steps_to_run(current_step, save_steps,
                                                steps_per_loop)
Hongkun Yu's avatar
Hongkun Yu committed
253
      train_steps(train_iterator, tf.convert_to_tensor(steps, dtype=tf.int32))
Hongkun Yu's avatar
Hongkun Yu committed
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
      current_step += steps
      train_loss = _float_metric_value(train_loss_metric)
      log_stream = "Train step: %d/%d  /  lr = %.9f  /  loss = %.7f" % (
          current_step, total_training_steps, learning_rate_fn(current_step),
          train_loss)
      if train_metric:
        log_stream += "  /  %s = %f" % (train_metric.name,
                                        _float_metric_value(train_metric))
      logging.info(log_stream)
      if train_summary_writer:
        with train_summary_writer.as_default():
          tf.summary.scalar(
              "learning_rate",
              learning_rate_fn(current_step),
              step=current_step)
          tf.summary.scalar(
              train_loss_metric.name, train_loss, step=current_step)
          if train_metric:
            tf.summary.scalar(
                train_metric.name,
                _float_metric_value(train_metric),
                step=current_step)
          train_summary_writer.flush()
277
278
279
      if model_dir and current_step % save_steps == 0:
        _save_checkpoint(checkpoint, model_dir,
                         checkpoint_name.format(step=current_step))
Hongkun Yu's avatar
Hongkun Yu committed
280

Hongkun Yu's avatar
Hongkun Yu committed
281
      if eval_fn and current_step % save_steps == 0:
Hongkun Yu's avatar
Hongkun Yu committed
282
283
284
285
286
287
288

        logging.info("Running evaluation after step: %s.", current_step)

        eval_fn(model, current_step, eval_summary_writer)
    if model_dir:
      _save_checkpoint(checkpoint, model_dir,
                       checkpoint_name.format(step=current_step))
Hongkun Yu's avatar
Hongkun Yu committed
289
    if eval_fn:
Hongkun Yu's avatar
Hongkun Yu committed
290
      logging.info("Running final evaluation after training is complete.")
291
292
293
294
295
296
297
298
      eval_metric = eval_fn(model, current_step, eval_summary_writer)

    training_summary = {
        "total_training_steps": total_training_steps,
        "train_loss": _float_metric_value(train_loss_metric),
    }
    if train_metric:
      training_summary["last_train_metrics"] = _float_metric_value(train_metric)
Hongkun Yu's avatar
Hongkun Yu committed
299
    if eval_fn:
300
301
302
      # eval_metric is supposed to be a float.
      training_summary["eval_metrics"] = eval_metric

Hongkun Yu's avatar
Hongkun Yu committed
303
    model_training_utils.write_txt_summary(training_summary, summary_dir)
Hongkun Yu's avatar
Hongkun Yu committed
304

Hongkun Yu's avatar
Hongkun Yu committed
305
    return model