train.py 4.13 KB
Newer Older
1
#!/usr/bin/env python3
2
3
import logging
import pathlib
4
from argparse import ArgumentParser
5

6
from common import MODEL_TYPE_LIBRISPEECH, MODEL_TYPE_MUSTC, MODEL_TYPE_TEDLIUM3
7
from librispeech.lightning import LibriSpeechRNNTModule
8
from mustc.lightning import MuSTCRNNTModule
9
10
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
11
from tedlium3.lightning import TEDLIUM3RNNTModule
12
13


14
def get_trainer(args):
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    checkpoint_dir = args.exp_dir / "checkpoints"
    checkpoint = ModelCheckpoint(
        checkpoint_dir,
        monitor="Losses/val_loss",
        mode="min",
        save_top_k=5,
        save_weights_only=True,
        verbose=True,
    )
    train_checkpoint = ModelCheckpoint(
        checkpoint_dir,
        monitor="Losses/train_loss",
        mode="min",
        save_top_k=5,
        save_weights_only=True,
        verbose=True,
    )
    callbacks = [
        checkpoint,
        train_checkpoint,
    ]
36
    return Trainer(
37
38
39
40
41
42
        default_root_dir=args.exp_dir,
        max_epochs=args.epochs,
        num_nodes=args.num_nodes,
        gpus=args.gpus,
        accelerator="gpu",
        strategy="ddp",
43
        gradient_clip_val=args.gradient_clip_val,
44
45
46
47
        callbacks=callbacks,
    )


48
49
50
51
52
53
54
55
56
57
58
59
60
def get_lightning_module(args):
    if args.model_type == MODEL_TYPE_LIBRISPEECH:
        return LibriSpeechRNNTModule(
            librispeech_path=str(args.dataset_path),
            sp_model_path=str(args.sp_model_path),
            global_stats_path=str(args.global_stats_path),
        )
    elif args.model_type == MODEL_TYPE_TEDLIUM3:
        return TEDLIUM3RNNTModule(
            tedlium_path=str(args.dataset_path),
            sp_model_path=str(args.sp_model_path),
            global_stats_path=str(args.global_stats_path),
        )
61
62
63
64
65
66
    elif args.model_type == MODEL_TYPE_MUSTC:
        return MuSTCRNNTModule(
            mustc_path=str(args.dataset_path),
            sp_model_path=str(args.sp_model_path),
            global_stats_path=str(args.global_stats_path),
        )
67
68
    else:
        raise ValueError(f"Encountered unsupported model type {args.model_type}.")
69

70
71
72

def parse_args():
    parser = ArgumentParser()
73
    parser.add_argument(
74
75
76
77
        "--model-type", type=str, choices=[MODEL_TYPE_LIBRISPEECH, MODEL_TYPE_TEDLIUM3, MODEL_TYPE_MUSTC], required=True
    )
    parser.add_argument(
        "--global-stats-path",
78
79
80
        default=pathlib.Path("global_stats.json"),
        type=pathlib.Path,
        help="Path to JSON file containing feature means and stddevs.",
81
        required=True,
82
83
    )
    parser.add_argument(
84
        "--dataset-path",
85
        type=pathlib.Path,
86
        help="Path to datasets.",
87
88
89
        required=True,
    )
    parser.add_argument(
90
        "--sp-model-path",
91
92
93
        type=pathlib.Path,
        help="Path to SentencePiece model.",
        required=True,
94
95
    )
    parser.add_argument(
96
        "--exp-dir",
97
        default=pathlib.Path("./exp"),
98
        type=pathlib.Path,
99
        help="Directory to save checkpoints and logs to. (Default: './exp')",
100
101
    )
    parser.add_argument(
102
        "--num-nodes",
103
        default=4,
104
        type=int,
105
        help="Number of nodes to use for training. (Default: 4)",
106
107
108
109
110
111
112
113
114
115
116
117
118
    )
    parser.add_argument(
        "--gpus",
        default=8,
        type=int,
        help="Number of GPUs per node to use for training. (Default: 8)",
    )
    parser.add_argument(
        "--epochs",
        default=120,
        type=int,
        help="Number of epochs to train for. (Default: 120)",
    )
119
    parser.add_argument(
120
        "--gradient-clip-val", default=10.0, type=float, help="Value to clip gradient values to. (Default: 10.0)"
121
    )
122
123
124
125
    parser.add_argument("--debug", action="store_true", help="whether to use debug level for logging")
    return parser.parse_args()


126
def init_logger(debug):
127
128
129
130
131
132
    fmt = "%(asctime)s %(message)s" if debug else "%(message)s"
    level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(format=fmt, level=level, datefmt="%Y-%m-%d %H:%M:%S")


def cli_main():
133
134
135
136
137
    args = parse_args()
    init_logger(args.debug)
    model = get_lightning_module(args)
    trainer = get_trainer(args)
    trainer.fit(model)
138
139
140
141


if __name__ == "__main__":
    cli_main()