train_lib.py 9.86 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 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
14
15
16
17
#
# 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.

"""Multitask training driver library."""
# pytype: disable=attribute-error
import os
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
18
from typing import Any, List, Mapping, Optional, Tuple, Union
Hongkun Yu's avatar
Hongkun Yu committed
19
20
21
22
23
from absl import logging
import orbit
import tensorflow as tf
from official.core import base_task
from official.core import base_trainer as core_lib
24
from official.core import train_utils
25
26
from official.modeling.multitask import base_model
from official.modeling.multitask import base_trainer
Hongkun Yu's avatar
Hongkun Yu committed
27
28
from official.modeling.multitask import configs
from official.modeling.multitask import evaluator as evaluator_lib
29
from official.modeling.multitask import interleaving_trainer
Hongkun Yu's avatar
Hongkun Yu committed
30
from official.modeling.multitask import multitask
31
32
33
34
35
36
37
38
from official.modeling.multitask import task_sampler

TRAINERS = {
    'interleaving': interleaving_trainer.MultiTaskInterleavingTrainer,
    'joint': base_trainer.MultiTaskBaseTrainer
}


Terry Huang's avatar
Terry Huang committed
39
40
41
42
43
44
45
46
def run_experiment(
    *,
    distribution_strategy: tf.distribute.Strategy,
    task: multitask.MultiTask,
    model: base_model.MultiTaskBaseModel,
    mode: str,
    params: configs.MultiTaskExperimentConfig,
    model_dir: str,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
47
    run_post_eval: bool = False,
Terry Huang's avatar
Terry Huang committed
48
    trainer: base_trainer.MultiTaskBaseTrainer = None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
49
50
    ) -> Union[base_model.MultiTaskBaseModel,
               Tuple[base_model.MultiTaskBaseModel, Mapping[Any, Any]]]:
51
52
53
54
55
56
57
58
59
60
  """Runs train/eval configured by the experiment params.

  Args:
    distribution_strategy: A distribution distribution_strategy.
    task: A MultiTaskTask instance.
    model: A MultiTaskBaseModel instance.
    mode: A 'str', specifying the mode. Can be 'train', 'eval', 'train_and_eval'
      or 'continuous_eval'.
    params: ExperimentConfig instance.
    model_dir: A 'str', a path to store model checkpoints and summaries.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
61
62
    run_post_eval: Whether to run post eval once after training, metrics logs
      are returned.
Terry Huang's avatar
Terry Huang committed
63
64
    trainer: (optional) A multi-task trainer to use. If none is provided, a
      default one will be created based on `params`.
65
66
67
68
69
70
71
72

  Returns:
      model: `base_model.MultiTaskBaseModel` instance.
  """

  is_training = 'train' in mode
  is_eval = 'eval' in mode
  with distribution_strategy.scope():
Frederick Liu's avatar
Frederick Liu committed
73
    optimizer = train_utils.create_optimizer(task, params)
74
75
76
77
78
    kwargs = dict(multi_task=task, multi_task_model=model, optimizer=optimizer)
    if params.trainer.trainer_type == 'interleaving':
      sampler = task_sampler.get_task_sampler(params.trainer.task_sampler,
                                              task.task_weights)
      kwargs.update(dict(task_sampler=sampler))
Terry Huang's avatar
Terry Huang committed
79
80
81
    if trainer is None:
      trainer = TRAINERS[params.trainer.trainer_type](
          **kwargs) if is_training else None
82
    if is_eval:
83
      eval_steps = task.task_eval_steps
84
      evaluator = evaluator_lib.MultiTaskEvaluator(
85
          eval_tasks=task.tasks.values(),
86
          model=model,
87
          eval_steps=eval_steps,
Tianqi Liu's avatar
Tianqi Liu committed
88
89
90
          global_step=trainer.global_step if is_training else None,
          checkpoint_exporter=train_utils.maybe_create_best_ckpt_exporter(
              params, model_dir))
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
    else:
      evaluator = None

  if trainer:
    checkpoint = trainer.checkpoint
    global_step = trainer.global_step
  else:
    checkpoint = evaluator.checkpoint
    global_step = evaluator.global_step

  checkpoint_manager = tf.train.CheckpointManager(
      checkpoint,
      directory=model_dir,
      max_to_keep=params.trainer.max_to_keep,
      step_counter=global_step,
      checkpoint_interval=params.trainer.checkpoint_interval,
      init_fn=model.initialize)

  controller = orbit.Controller(
      strategy=distribution_strategy,
      trainer=trainer,
      evaluator=evaluator,
      global_step=global_step,
      steps_per_loop=params.trainer.steps_per_loop,
      checkpoint_manager=checkpoint_manager,
      summary_dir=os.path.join(model_dir, 'train'),
      eval_summary_dir=os.path.join(model_dir, 'validation'),
      summary_interval=params.trainer.summary_interval)

  logging.info('Starts to execute mode: %s', mode)
  with distribution_strategy.scope():
    if mode == 'train':
      controller.train(steps=params.trainer.train_steps)
    elif mode == 'train_and_eval':
      controller.train_and_evaluate(
          train_steps=params.trainer.train_steps,
          eval_steps=params.trainer.validation_steps,
          eval_interval=params.trainer.validation_interval)
    elif mode == 'eval':
      controller.evaluate(steps=params.trainer.validation_steps)
    elif mode == 'continuous_eval':

      def timeout_fn():
        if evaluator.global_step.numpy() >= params.trainer.train_steps:
          return True
        return False

      controller.evaluate_continuously(
          steps=params.trainer.validation_steps,
          timeout=params.trainer.continuous_eval_timeout,
          timeout_fn=timeout_fn)
    else:
      raise NotImplementedError('The mode is not implemented: %s' % mode)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
145
146
147
148
149
    if run_post_eval:
      return model, evaluator.evaluate(
          tf.convert_to_tensor(params.trainer.validation_steps))  # pytype: disable=bad-return-type  # typed-keras
    else:
      return model
Hongkun Yu's avatar
Hongkun Yu committed
150
151


152
def run_experiment_with_multitask_eval(
Hongkun Yu's avatar
Hongkun Yu committed
153
    *,
Hongkun Yu's avatar
Hongkun Yu committed
154
155
    distribution_strategy: tf.distribute.Strategy,
    train_task: base_task.Task,
156
    eval_tasks: List[base_task.Task],
Hongkun Yu's avatar
Hongkun Yu committed
157
    mode: str,
Hongkun Yu's avatar
Hongkun Yu committed
158
    params: configs.MultiEvalExperimentConfig,
Hongkun Yu's avatar
Hongkun Yu committed
159
160
    model_dir: str,
    run_post_eval: bool = False,
Le Hou's avatar
Le Hou committed
161
    save_summary: bool = True,
Terry Huang's avatar
Terry Huang committed
162
    trainer: Optional[core_lib.Trainer] = None) -> Tuple[Any, Any]:
Hongkun Yu's avatar
Hongkun Yu committed
163
164
165
166
167
  """Runs train/eval configured by the experiment params.

  Args:
    distribution_strategy: A distribution distribution_strategy.
    train_task: A base_task.Task instance.
168
    eval_tasks: A list of evaluation tasks.
Hongkun Yu's avatar
Hongkun Yu committed
169
170
171
172
    mode: A 'str', specifying the mode. Can be 'train', 'eval', 'train_and_eval'
      or 'continuous_eval'.
    params: MultiEvalExperimentConfig instance.
    model_dir: A 'str', a path to store model checkpoints and summaries.
Hongkun Yu's avatar
Hongkun Yu committed
173
174
175
    run_post_eval: Whether to run post eval once after training, metrics logs
      are returned.
    save_summary: Whether to save train and validation summary.
Le Hou's avatar
Le Hou committed
176
177
178
    trainer: the core_lib.Trainer instance. It should be created within the
      strategy.scope(). If not provided, an instance will be created by default
      if `mode` contains 'train'.
Hongkun Yu's avatar
Hongkun Yu committed
179
180
181
182
183
184
185
186
187

  Returns:
      model: `tf.keras.Model` instance.
  """

  is_training = 'train' in mode
  is_eval = 'eval' in mode
  with distribution_strategy.scope():
    if is_training:
Le Hou's avatar
Le Hou committed
188
      trainer = trainer or core_lib.Trainer(
Hongkun Yu's avatar
Hongkun Yu committed
189
190
          config=params,
          task=train_task,
Le Hou's avatar
Le Hou committed
191
          model=train_task.build_model(),
Frederick Liu's avatar
Frederick Liu committed
192
          optimizer=train_utils.create_optimizer(train_task, params),
Hongkun Yu's avatar
Hongkun Yu committed
193
194
195
196
          train=True,
          evaluate=False)
    else:
      trainer = None
Le Hou's avatar
Le Hou committed
197
198
    model = trainer.model if trainer else train_task.build_model()

Hongkun Yu's avatar
Hongkun Yu committed
199
    if is_eval:
200
201
202
      eval_steps = dict([(task_routine.task_config.name,
                          task_routine.eval_steps)
                         for task_routine in params.eval_tasks])
Hongkun Yu's avatar
Hongkun Yu committed
203
      evaluator = evaluator_lib.MultiTaskEvaluator(
204
          eval_tasks=eval_tasks,
Hongkun Yu's avatar
Hongkun Yu committed
205
          model=model,
206
          global_step=trainer.global_step if is_training else None,
207
          eval_steps=eval_steps,
208
209
          checkpoint_exporter=train_utils.maybe_create_best_ckpt_exporter(
              params, model_dir))
Hongkun Yu's avatar
Hongkun Yu committed
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
    else:
      evaluator = None

  if trainer:
    checkpoint = trainer.checkpoint
    global_step = trainer.global_step
  else:
    checkpoint = evaluator.checkpoint
    global_step = evaluator.global_step

  checkpoint_manager = tf.train.CheckpointManager(
      checkpoint,
      directory=model_dir,
      max_to_keep=params.trainer.max_to_keep,
      step_counter=global_step,
      checkpoint_interval=params.trainer.checkpoint_interval,
      init_fn=trainer.initialize if trainer else None)

  controller = orbit.Controller(
      strategy=distribution_strategy,
      trainer=trainer,
      evaluator=evaluator,
      global_step=global_step,
      steps_per_loop=params.trainer.steps_per_loop,
      checkpoint_manager=checkpoint_manager,
Hongkun Yu's avatar
Hongkun Yu committed
235
236
237
238
239
      summary_dir=os.path.join(model_dir, 'train') if save_summary else None,
      eval_summary_dir=os.path.join(model_dir, 'validation') if
      (save_summary) else None,
      summary_interval=params.trainer.summary_interval if
      (save_summary) else None)
Hongkun Yu's avatar
Hongkun Yu committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265

  logging.info('Starts to execute mode: %s', mode)
  with distribution_strategy.scope():
    if mode == 'train':
      controller.train(steps=params.trainer.train_steps)
    elif mode == 'train_and_eval':
      controller.train_and_evaluate(
          train_steps=params.trainer.train_steps,
          eval_steps=params.trainer.validation_steps,
          eval_interval=params.trainer.validation_interval)
    elif mode == 'eval':
      controller.evaluate(steps=params.trainer.validation_steps)
    elif mode == 'continuous_eval':

      def timeout_fn():
        if evaluator.global_step.numpy() >= params.trainer.train_steps:
          return True
        return False

      controller.evaluate_continuously(
          steps=params.trainer.validation_steps,
          timeout=params.trainer.continuous_eval_timeout,
          timeout_fn=timeout_fn)
    else:
      raise NotImplementedError('The mode is not implemented: %s' % mode)

Hongkun Yu's avatar
Hongkun Yu committed
266
267
    if run_post_eval:
      return model, evaluator.evaluate(
Rebecca Chen's avatar
Rebecca Chen committed
268
          tf.convert_to_tensor(params.trainer.validation_steps))  # pytype: disable=bad-return-type  # typed-keras
Hongkun Yu's avatar
Hongkun Yu committed
269
    else:
Rebecca Chen's avatar
Rebecca Chen committed
270
      return model, {}  # pytype: disable=bad-return-type  # typed-keras