File size: 13,059 Bytes
0d80816
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
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
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
185
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
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
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
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import collections
import json
import os
import sys
import time

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import ConcatDataset, DataLoader
from torch.utils.tensorboard import SummaryWriter

from models.base.base_sampler import BatchSampler
from utils.util import (
    Logger,
    remove_older_ckpt,
    save_config,
    set_all_random_seed,
    ValueWindow,
)


class BaseTrainer(object):
    def __init__(self, args, cfg):
        self.args = args
        self.log_dir = args.log_dir
        self.cfg = cfg

        self.checkpoint_dir = os.path.join(args.log_dir, "checkpoints")
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        if not cfg.train.ddp or args.local_rank == 0:
            self.sw = SummaryWriter(os.path.join(args.log_dir, "events"))
            self.logger = self.build_logger()
        self.time_window = ValueWindow(50)

        self.step = 0
        self.epoch = -1
        self.max_epochs = self.cfg.train.epochs
        self.max_steps = self.cfg.train.max_steps

        # set random seed & init distributed training
        set_all_random_seed(self.cfg.train.random_seed)
        if cfg.train.ddp:
            dist.init_process_group(backend="nccl")

        if cfg.model_type not in ["AutoencoderKL", "AudioLDM"]:
            self.singers = self.build_singers_lut()

        # setup data_loader
        self.data_loader = self.build_data_loader()

        # setup model & enable distributed training
        self.model = self.build_model()
        print(self.model)

        if isinstance(self.model, dict):
            for key, value in self.model.items():
                value.cuda(self.args.local_rank)
                if key == "PQMF":
                    continue
                if cfg.train.ddp:
                    self.model[key] = DistributedDataParallel(
                        value, device_ids=[self.args.local_rank]
                    )
        else:
            self.model.cuda(self.args.local_rank)
            if cfg.train.ddp:
                self.model = DistributedDataParallel(
                    self.model, device_ids=[self.args.local_rank]
                )

        # create criterion
        self.criterion = self.build_criterion()
        if isinstance(self.criterion, dict):
            for key, value in self.criterion.items():
                self.criterion[key].cuda(args.local_rank)
        else:
            self.criterion.cuda(self.args.local_rank)

        # optimizer
        self.optimizer = self.build_optimizer()
        self.scheduler = self.build_scheduler()

        # save config file
        self.config_save_path = os.path.join(self.checkpoint_dir, "args.json")

    def build_logger(self):
        log_file = os.path.join(self.checkpoint_dir, "train.log")
        logger = Logger(log_file, level=self.args.log_level).logger

        return logger

    def build_dataset(self):
        raise NotImplementedError

    def build_data_loader(self):
        Dataset, Collator = self.build_dataset()
        # build dataset instance for each dataset and combine them by ConcatDataset
        datasets_list = []
        for dataset in self.cfg.dataset:
            subdataset = Dataset(self.cfg, dataset, is_valid=False)
            datasets_list.append(subdataset)
        train_dataset = ConcatDataset(datasets_list)

        train_collate = Collator(self.cfg)
        # TODO: multi-GPU training
        if self.cfg.train.ddp:
            raise NotImplementedError("DDP is not supported yet.")

        # sampler will provide indices to batch_sampler, which will perform batching and yield batch indices
        batch_sampler = BatchSampler(
            cfg=self.cfg, concat_dataset=train_dataset, dataset_list=datasets_list
        )

        # use batch_sampler argument instead of (sampler, shuffle, drop_last, batch_size)
        train_loader = DataLoader(
            train_dataset,
            collate_fn=train_collate,
            num_workers=self.args.num_workers,
            batch_sampler=batch_sampler,
            pin_memory=False,
        )
        if not self.cfg.train.ddp or self.args.local_rank == 0:
            datasets_list = []
            for dataset in self.cfg.dataset:
                subdataset = Dataset(self.cfg, dataset, is_valid=True)
                datasets_list.append(subdataset)
            valid_dataset = ConcatDataset(datasets_list)
            valid_collate = Collator(self.cfg)
            batch_sampler = BatchSampler(
                cfg=self.cfg, concat_dataset=valid_dataset, dataset_list=datasets_list
            )
            valid_loader = DataLoader(
                valid_dataset,
                collate_fn=valid_collate,
                num_workers=1,
                batch_sampler=batch_sampler,
            )
        else:
            raise NotImplementedError("DDP is not supported yet.")
            # valid_loader = None
        data_loader = {"train": train_loader, "valid": valid_loader}
        return data_loader

    def build_singers_lut(self):
        # combine singers
        if not os.path.exists(os.path.join(self.log_dir, self.cfg.preprocess.spk2id)):
            singers = collections.OrderedDict()
        else:
            with open(
                os.path.join(self.log_dir, self.cfg.preprocess.spk2id), "r"
            ) as singer_file:
                singers = json.load(singer_file)
        singer_count = len(singers)
        for dataset in self.cfg.dataset:
            singer_lut_path = os.path.join(
                self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.spk2id
            )
            with open(singer_lut_path, "r") as singer_lut_path:
                singer_lut = json.load(singer_lut_path)
            for singer in singer_lut.keys():
                if singer not in singers:
                    singers[singer] = singer_count
                    singer_count += 1
        with open(
            os.path.join(self.log_dir, self.cfg.preprocess.spk2id), "w"
        ) as singer_file:
            json.dump(singers, singer_file, indent=4, ensure_ascii=False)
        print(
            "singers have been dumped to {}".format(
                os.path.join(self.log_dir, self.cfg.preprocess.spk2id)
            )
        )
        return singers

    def build_model(self):
        raise NotImplementedError()

    def build_optimizer(self):
        raise NotImplementedError

    def build_scheduler(self):
        raise NotImplementedError()

    def build_criterion(self):
        raise NotImplementedError

    def get_state_dict(self):
        raise NotImplementedError

    def save_config_file(self):
        save_config(self.config_save_path, self.cfg)

    # TODO, save without module.
    def save_checkpoint(self, state_dict, saved_model_path):
        torch.save(state_dict, saved_model_path)

    def load_checkpoint(self):
        checkpoint_path = os.path.join(self.checkpoint_dir, "checkpoint")
        assert os.path.exists(checkpoint_path)
        checkpoint_filename = open(checkpoint_path).readlines()[-1].strip()
        model_path = os.path.join(self.checkpoint_dir, checkpoint_filename)
        assert os.path.exists(model_path)
        if not self.cfg.train.ddp or self.args.local_rank == 0:
            self.logger.info(f"Re(store) from {model_path}")
        checkpoint = torch.load(model_path, map_location="cpu")
        return checkpoint

    def load_model(self, checkpoint):
        raise NotImplementedError

    def restore(self):
        checkpoint = self.load_checkpoint()
        self.load_model(checkpoint)

    def train_step(self, data):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    @torch.no_grad()
    def eval_step(self):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    def write_summary(self, losses, stats):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    def write_valid_summary(self, losses, stats):
        raise NotImplementedError(
            f"Need to implement function {sys._getframe().f_code.co_name} in "
            f"your sub-class of {self.__class__.__name__}. "
        )

    def echo_log(self, losses, mode="Training"):
        message = [
            "{} - Epoch {} Step {}: [{:.3f} s/step]".format(
                mode, self.epoch + 1, self.step, self.time_window.average
            )
        ]

        for key in sorted(losses.keys()):
            if isinstance(losses[key], dict):
                for k, v in losses[key].items():
                    message.append(
                        str(k).split("/")[-1] + "=" + str(round(float(v), 5))
                    )
            else:
                message.append(
                    str(key).split("/")[-1] + "=" + str(round(float(losses[key]), 5))
                )
        self.logger.info(", ".join(message))

    def eval_epoch(self):
        self.logger.info("Validation...")
        valid_losses = {}
        for i, batch_data in enumerate(self.data_loader["valid"]):
            for k, v in batch_data.items():
                if isinstance(v, torch.Tensor):
                    batch_data[k] = v.cuda()
            valid_loss, valid_stats, total_valid_loss = self.eval_step(batch_data, i)
            for key in valid_loss:
                if key not in valid_losses:
                    valid_losses[key] = 0
                valid_losses[key] += valid_loss[key]

        # Add mel and audio to the Tensorboard
        # Average loss
        for key in valid_losses:
            valid_losses[key] /= i + 1
        self.echo_log(valid_losses, "Valid")
        return valid_losses, valid_stats

    def train_epoch(self):
        for i, batch_data in enumerate(self.data_loader["train"]):
            start_time = time.time()
            # Put the data to cuda device
            for k, v in batch_data.items():
                if isinstance(v, torch.Tensor):
                    batch_data[k] = v.cuda(self.args.local_rank)

            # Training step
            train_losses, train_stats, total_loss = self.train_step(batch_data)
            self.time_window.append(time.time() - start_time)

            if self.args.local_rank == 0 or not self.cfg.train.ddp:
                if self.step % self.args.stdout_interval == 0:
                    self.echo_log(train_losses, "Training")

                if self.step % self.cfg.train.save_summary_steps == 0:
                    self.logger.info(f"Save summary as step {self.step}")
                    self.write_summary(train_losses, train_stats)

                if (
                    self.step % self.cfg.train.save_checkpoints_steps == 0
                    and self.step != 0
                ):
                    saved_model_name = "step-{:07d}_loss-{:.4f}.pt".format(
                        self.step, total_loss
                    )
                    saved_model_path = os.path.join(
                        self.checkpoint_dir, saved_model_name
                    )
                    saved_state_dict = self.get_state_dict()
                    self.save_checkpoint(saved_state_dict, saved_model_path)
                    self.save_config_file()
                    # keep max n models
                    remove_older_ckpt(
                        saved_model_name,
                        self.checkpoint_dir,
                        max_to_keep=self.cfg.train.keep_checkpoint_max,
                    )

                if self.step != 0 and self.step % self.cfg.train.valid_interval == 0:
                    if isinstance(self.model, dict):
                        for key in self.model.keys():
                            self.model[key].eval()
                    else:
                        self.model.eval()
                    # Evaluate one epoch and get average loss
                    valid_losses, valid_stats = self.eval_epoch()
                    if isinstance(self.model, dict):
                        for key in self.model.keys():
                            self.model[key].train()
                    else:
                        self.model.train()
                    # Write validation losses to summary.
                    self.write_valid_summary(valid_losses, valid_stats)
            self.step += 1

    def train(self):
        for epoch in range(max(0, self.epoch), self.max_epochs):
            self.train_epoch()
            self.epoch += 1
            if self.step > self.max_steps:
                self.logger.info("Training finished!")
                break