lightning_base.py 14.8 KB
Newer Older
1
import argparse
2
import logging
3
import os
4
5
from pathlib import Path
from typing import Any, Dict
6
7

import pytorch_lightning as pl
8
from pytorch_lightning.utilities import rank_zero_info
9

10
import pkg_resources
11
12
from transformers import (
    AdamW,
13
14
15
16
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
17
    AutoModelForSeq2SeqLM,
18
19
20
21
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
22
23
    PretrainedConfig,
    PreTrainedTokenizer,
24
25
)
from transformers.optimization import (
Sam Shleifer's avatar
Sam Shleifer committed
26
    Adafactor,
27
28
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
29
    get_linear_schedule_with_warmup,
30
    get_polynomial_decay_schedule_with_warmup,
31
32
33
)


34
35
logger = logging.getLogger(__name__)

36
37
38
39
40
41
42
43
44
try:
    pkg = "pytorch_lightning"
    min_ver = "1.0.4"
    pkg_resources.require(f"{pkg}>={min_ver}")
except pkg_resources.VersionConflict:
    logger.warning(
        f"{pkg}>={min_ver} is required for a normal functioning of this module, but found {pkg}=={pkg_resources.get_distribution(pkg).version}. Try pip install -r examples/requirements.txt"
    )

45

46
47
48
49
50
51
52
MODEL_MODES = {
    "base": AutoModel,
    "sequence-classification": AutoModelForSequenceClassification,
    "question-answering": AutoModelForQuestionAnswering,
    "pretraining": AutoModelForPreTraining,
    "token-classification": AutoModelForTokenClassification,
    "language-modeling": AutoModelWithLMHead,
53
54
    "summarization": AutoModelForSeq2SeqLM,
    "translation": AutoModelForSeq2SeqLM,
55
56
57
}


58
59
60
61
62
# update this and the import above to support new schedulers from transformers.optimization
arg_to_scheduler = {
    "linear": get_linear_schedule_with_warmup,
    "cosine": get_cosine_schedule_with_warmup,
    "cosine_w_restarts": get_cosine_with_hard_restarts_schedule_with_warmup,
63
    "polynomial": get_polynomial_decay_schedule_with_warmup,
64
65
66
67
68
69
70
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
arg_to_scheduler_choices = sorted(arg_to_scheduler.keys())
arg_to_scheduler_metavar = "{" + ", ".join(arg_to_scheduler_choices) + "}"


71
class BaseTransformer(pl.LightningModule):
72
73
74
75
76
77
78
79
80
81
    def __init__(
        self,
        hparams: argparse.Namespace,
        num_labels=None,
        mode="base",
        config=None,
        tokenizer=None,
        model=None,
        **config_kwargs
    ):
82
        """Initialize a model, tokenizer and config."""
Julien Chaumond's avatar
Julien Chaumond committed
83
        super().__init__()
84
85
86
87
        # TODO: move to self.save_hyperparameters()
        # self.save_hyperparameters()
        # can also expand arguments into trainer signature for easier reading

88
        self.save_hyperparameters(hparams)
89
90
        self.step_count = 0
        self.output_dir = Path(self.hparams.output_dir)
91
        cache_dir = self.hparams.cache_dir if self.hparams.cache_dir else None
92
93
94
95
96
97
98
99
100
        if config is None:
            self.config = AutoConfig.from_pretrained(
                self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path,
                **({"num_labels": num_labels} if num_labels is not None else {}),
                cache_dir=cache_dir,
                **config_kwargs,
            )
        else:
            self.config: PretrainedConfig = config
101
102
103
104
105
106
107

        extra_model_params = ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
        for p in extra_model_params:
            if getattr(self.hparams, p, None):
                assert hasattr(self.config, p), f"model config doesn't have a `{p}` attribute"
                setattr(self.config, p, getattr(self.hparams, p))

108
109
110
111
112
113
114
        if tokenizer is None:
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path,
                cache_dir=cache_dir,
            )
        else:
            self.tokenizer: PreTrainedTokenizer = tokenizer
115
        self.model_type = MODEL_MODES[mode]
116
117
118
119
120
121
122
123
124
125
126
127
        if model is None:
            self.model = self.model_type.from_pretrained(
                self.hparams.model_name_or_path,
                from_tf=bool(".ckpt" in self.hparams.model_name_or_path),
                config=self.config,
                cache_dir=cache_dir,
            )
        else:
            self.model = model

    def load_hf_checkpoint(self, *args, **kwargs):
        self.model = self.model_type.from_pretrained(*args, **kwargs)
128

129
130
131
    def get_lr_scheduler(self):
        get_schedule_func = arg_to_scheduler[self.hparams.lr_scheduler]
        scheduler = get_schedule_func(
132
            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps()
133
134
135
136
        )
        scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
        return scheduler

137
    def configure_optimizers(self):
138
        """Prepare optimizer and schedule (linear warmup and decay)"""
139
        model = self.model
140
141
142
143
144
145
146
147
148
149
150
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": self.hparams.weight_decay,
            },
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
            },
        ]
Sam Shleifer's avatar
Sam Shleifer committed
151
152
153
154
155
156
157
158
159
        if self.hparams.adafactor:
            optimizer = Adafactor(
                optimizer_grouped_parameters, lr=self.hparams.learning_rate, scale_parameter=False, relative_step=False
            )

        else:
            optimizer = AdamW(
                optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon
            )
160
        self.opt = optimizer
161

162
163
        scheduler = self.get_lr_scheduler()

164
        return [optimizer], [scheduler]
165

166
167
168
    def test_step(self, batch, batch_nb):
        return self.validation_step(batch, batch_nb)

169
    def test_epoch_end(self, outputs):
170
171
        return self.validation_end(outputs)

172
173
174
175
    def total_steps(self) -> int:
        """The number of total training steps that will be run. Used for lr scheduler purposes."""
        num_devices = max(1, self.hparams.gpus)  # TODO: consider num_tpu_cores
        effective_batch_size = self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
176
        return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
177
178

    def setup(self, mode):
179
180
181
        if mode == "test":
            self.dataset_size = len(self.test_dataloader().dataset)
        else:
182
            self.train_loader = self.get_dataloader("train", self.hparams.train_batch_size, shuffle=True)
183
            self.dataset_size = len(self.train_dataloader().dataset)
184

185
    def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False):
186
        raise NotImplementedError("You must implement this for your task")
187
188
189

    def train_dataloader(self):
        return self.train_loader
190
191

    def val_dataloader(self):
192
        return self.get_dataloader("dev", self.hparams.eval_batch_size, shuffle=False)
193
194

    def test_dataloader(self):
195
        return self.get_dataloader("test", self.hparams.eval_batch_size, shuffle=False)
196

197
198
199
200
201
202
203
204
205
206
    def _feature_file(self, mode):
        return os.path.join(
            self.hparams.data_dir,
            "cached_{}_{}_{}".format(
                mode,
                list(filter(None, self.hparams.model_name_or_path.split("/"))).pop(),
                str(self.hparams.max_seq_length),
            ),
        )

207
208
209
210
211
212
213
    @pl.utilities.rank_zero_only
    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        save_path = self.output_dir.joinpath("best_tfmr")
        self.model.config.save_step = self.step_count
        self.model.save_pretrained(save_path)
        self.tokenizer.save_pretrained(save_path)

214
215
216
217
218
219
220
    @staticmethod
    def add_model_specific_args(parser, root_dir):
        parser.add_argument(
            "--model_name_or_path",
            default=None,
            type=str,
            required=True,
Julien Chaumond's avatar
Julien Chaumond committed
221
            help="Path to pretrained model or model identifier from huggingface.co/models",
222
223
224
225
226
227
        )
        parser.add_argument(
            "--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name"
        )
        parser.add_argument(
            "--tokenizer_name",
228
            default=None,
229
230
231
232
233
234
235
236
237
            type=str,
            help="Pretrained tokenizer name or path if not the same as model_name",
        )
        parser.add_argument(
            "--cache_dir",
            default="",
            type=str,
            help="Where do you want to store the pre-trained models downloaded from s3",
        )
238
239
240
241
242
243
244
245
246
247
248
        parser.add_argument(
            "--encoder_layerdrop",
            type=float,
            help="Encoder layer dropout probability (Optional). Goes into model.config",
        )
        parser.add_argument(
            "--decoder_layerdrop",
            type=float,
            help="Decoder layer dropout probability (Optional). Goes into model.config",
        )
        parser.add_argument(
Lysandre's avatar
Lysandre committed
249
250
251
            "--dropout",
            type=float,
            help="Dropout probability (Optional). Goes into model.config",
252
253
        )
        parser.add_argument(
Lysandre's avatar
Lysandre committed
254
255
256
            "--attention_dropout",
            type=float,
            help="Attention dropout probability (Optional). Goes into model.config",
257
        )
258
        parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
259
260
261
262
263
264
265
266
        parser.add_argument(
            "--lr_scheduler",
            default="linear",
            choices=arg_to_scheduler_choices,
            metavar=arg_to_scheduler_metavar,
            type=str,
            help="Learning rate scheduler",
        )
267
268
        parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.")
        parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
269
        parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
270
        parser.add_argument("--num_workers", default=4, type=int, help="kwarg passed to DataLoader")
271
        parser.add_argument("--num_train_epochs", dest="max_epochs", default=3, type=int)
272
273
        parser.add_argument("--train_batch_size", default=32, type=int)
        parser.add_argument("--eval_batch_size", default=32, type=int)
Sam Shleifer's avatar
Sam Shleifer committed
274
        parser.add_argument("--adafactor", action="store_true")
275
276


277
class LoggingCallback(pl.Callback):
278
    def on_batch_end(self, trainer, pl_module):
279
280
        lr_scheduler = trainer.lr_schedulers[0]["scheduler"]
        lrs = {f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr())}
281
282
        pl_module.logger.log_metrics(lrs)

283
    def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
284
285
286
287
288
289
290
291
        rank_zero_info("***** Validation results *****")
        metrics = trainer.callback_metrics
        # Log results
        for key in sorted(metrics):
            if key not in ["log", "progress_bar"]:
                rank_zero_info("{} = {}\n".format(key, str(metrics[key])))

    def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
292
        rank_zero_info("***** Test results *****")
293
294
295
296
        metrics = trainer.callback_metrics
        # Log and save results to file
        output_test_results_file = os.path.join(pl_module.hparams.output_dir, "test_results.txt")
        with open(output_test_results_file, "w") as writer:
297
298
            for key in sorted(metrics):
                if key not in ["log", "progress_bar"]:
299
                    rank_zero_info("{} = {}\n".format(key, str(metrics[key])))
300
                    writer.write("{} = {}\n".format(key, str(metrics[key])))
301
302


303
def add_generic_args(parser, root_dir) -> None:
304
305
    #  To allow all pl args uncomment the following line
    #  parser = pl.Trainer.add_argparse_args(parser)
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--fp16",
        action="store_true",
        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
    )

    parser.add_argument(
        "--fp16_opt_level",
        type=str,
322
        default="O2",
323
324
325
        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
        "See details at https://nvidia.github.io/apex/amp.html",
    )
326
    parser.add_argument("--n_tpu_cores", dest="tpu_cores", type=int)
327
    parser.add_argument("--max_grad_norm", dest="gradient_clip_val", default=1.0, type=float, help="Max gradient norm")
328
329
330
331
    parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
    parser.add_argument("--do_predict", action="store_true", help="Whether to run predictions on the test set.")
    parser.add_argument(
        "--gradient_accumulation_steps",
332
        dest="accumulate_grad_batches",
333
334
335
336
337
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
338
339
340
341
342
343
344
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input data dir. Should contain the training files for the CoNLL-2003 NER task.",
    )
345
346
347
348
349


def generic_train(
    model: BaseTransformer,
    args: argparse.Namespace,
350
    early_stopping_callback=None,
351
352
353
354
355
356
    logger=True,  # can pass WandbLogger() here
    extra_callbacks=[],
    checkpoint_callback=None,
    logging_callback=None,
    **extra_train_kwargs
):
357
358
    pl.seed_everything(args.seed)

359
    # init model
360
361
    odir = Path(model.hparams.output_dir)
    odir.mkdir(exist_ok=True)
362
363

    # add custom checkpoints
364
365
366
367
    if checkpoint_callback is None:
        checkpoint_callback = pl.callbacks.ModelCheckpoint(
            filepath=args.output_dir, prefix="checkpoint", monitor="val_loss", mode="min", save_top_k=1
        )
368
369
    if early_stopping_callback:
        extra_callbacks.append(early_stopping_callback)
370
371
    if logging_callback is None:
        logging_callback = LoggingCallback()
372

373
    train_params = {}
374

375
    # TODO: remove with PyTorch 1.6 since pl uses native amp
srush's avatar
srush committed
376
    if args.fp16:
377
        train_params["precision"] = 16
srush's avatar
srush committed
378
379
        train_params["amp_level"] = args.fp16_opt_level

380
    if args.gpus > 1:
srush's avatar
srush committed
381
382
        train_params["distributed_backend"] = "ddp"

Ola Piktus's avatar
Ola Piktus committed
383
384
    train_params["accumulate_grad_batches"] = args.accumulate_grad_batches

385
386
387
388
    trainer = pl.Trainer.from_argparse_args(
        args,
        weights_summary=None,
        callbacks=[logging_callback] + extra_callbacks,
389
390
391
392
        logger=logger,
        checkpoint_callback=checkpoint_callback,
        **train_params,
    )
srush's avatar
srush committed
393

394
395
    if args.do_train:
        trainer.fit(model)
396

397
    return trainer