train.py 6.6 KB
Newer Older
L
lifuchen 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# 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.
L
lifuchen 已提交
14 15 16 17 18 19
import numpy as np
import argparse
import os
import time
import math
from pathlib import Path
L
lifuchen 已提交
20 21
from parse import add_config_options_to_parser
from pprint import pprint
L
lifuchen 已提交
22
from ruamel import yaml
L
lifuchen 已提交
23
from tqdm import tqdm
L
lifuchen 已提交
24
from collections import OrderedDict
L
lifuchen 已提交
25 26 27 28
from tensorboardX import SummaryWriter
import paddle.fluid.dygraph as dg
import paddle.fluid.layers as layers
import paddle.fluid as fluid
L
lifuchen 已提交
29
from parakeet.models.transformer_tts.transformer_tts import TransformerTTS
L
lifuchen 已提交
30 31
from parakeet.models.fastspeech.fastspeech import FastSpeech
from parakeet.models.fastspeech.utils import get_alignment
L
lifuchen 已提交
32 33 34
import sys
sys.path.append("../transformer_tts")
from data import LJSpeechLoader
L
lifuchen 已提交
35

L
lifuchen 已提交
36

L
lifuchen 已提交
37
def load_checkpoint(step, model_path):
L
lifuchen 已提交
38 39
    model_dict, opti_dict = fluid.dygraph.load_dygraph(
        os.path.join(model_path, step))
L
lifuchen 已提交
40 41 42 43
    new_state_dict = OrderedDict()
    for param in model_dict:
        if param.startswith('_layers.'):
            new_state_dict[param[8:]] = model_dict[param]
L
lifuchen 已提交
44
        else:
L
lifuchen 已提交
45 46
            new_state_dict[param] = model_dict[param]
    return new_state_dict, opti_dict
L
lifuchen 已提交
47

L
lifuchen 已提交
48

L
lifuchen 已提交
49 50 51
def main(args):
    local_rank = dg.parallel.Env().local_rank if args.use_data_parallel else 0
    nranks = dg.parallel.Env().nranks if args.use_data_parallel else 1
L
lifuchen 已提交
52

L
lifuchen 已提交
53 54
    with open(args.config_path) as f:
        cfg = yaml.load(f, Loader=yaml.Loader)
L
lifuchen 已提交
55 56 57

    global_step = 0
    place = (fluid.CUDAPlace(dg.parallel.Env().dev_id)
L
lifuchen 已提交
58 59
             if args.use_data_parallel else fluid.CUDAPlace(0)
             if args.use_gpu else fluid.CPUPlace())
L
lifuchen 已提交
60

L
lifuchen 已提交
61
    if not os.path.exists(args.log_dir):
L
lifuchen 已提交
62 63
        os.mkdir(args.log_dir)
    path = os.path.join(args.log_dir, 'fastspeech')
L
lifuchen 已提交
64 65 66 67

    writer = SummaryWriter(path) if local_rank == 0 else None

    with dg.guard(place):
L
lifuchen 已提交
68 69
        with fluid.unique_name.guard():
            transformerTTS = TransformerTTS(cfg)
L
lifuchen 已提交
70 71 72
            model_dict, _ = load_checkpoint(
                str(args.transformer_step),
                os.path.join(args.transtts_path, "transformer"))
L
lifuchen 已提交
73 74
            transformerTTS.set_dict(model_dict)
            transformerTTS.eval()
L
lifuchen 已提交
75 76 77

        model = FastSpeech(cfg)
        model.train()
L
lifuchen 已提交
78 79 80 81 82 83 84
        optimizer = fluid.optimizer.AdamOptimizer(
            learning_rate=dg.NoamDecay(1 / (
                cfg['warm_up_step'] * (args.lr**2)), cfg['warm_up_step']),
            parameter_list=model.parameters())
        reader = LJSpeechLoader(
            cfg, args, nranks, local_rank, shuffle=True).reader()

L
lifuchen 已提交
85
        if args.checkpoint_path is not None:
L
lifuchen 已提交
86 87 88
            model_dict, opti_dict = load_checkpoint(
                str(args.fastspeech_step),
                os.path.join(args.checkpoint_path, "fastspeech"))
L
lifuchen 已提交
89 90
            model.set_dict(model_dict)
            optimizer.set_dict(opti_dict)
L
lifuchen 已提交
91
            global_step = args.fastspeech_step
L
lifuchen 已提交
92 93
            print("load checkpoint!!!")

L
lifuchen 已提交
94
        if args.use_data_parallel:
L
lifuchen 已提交
95
            strategy = dg.parallel.prepare_context()
L
lifuchen 已提交
96
            model = fluid.dygraph.parallel.DataParallel(model, strategy)
L
lifuchen 已提交
97

L
lifuchen 已提交
98
        for epoch in range(args.epochs):
L
lifuchen 已提交
99 100 101
            pbar = tqdm(reader)

            for i, data in enumerate(pbar):
L
lifuchen 已提交
102
                pbar.set_description('Processing at epoch %d' % epoch)
L
lifuchen 已提交
103
                character, mel, mel_input, pos_text, pos_mel, text_length, mel_lens = data
L
lifuchen 已提交
104

L
lifuchen 已提交
105 106 107 108 109
                _, _, attn_probs, _, _, _ = transformerTTS(
                    character, mel_input, pos_text, pos_mel)
                alignment = dg.to_variable(
                    get_alignment(attn_probs, mel_lens, cfg[
                        'transformer_head'])).astype(np.float32)
L
lifuchen 已提交
110

L
lifuchen 已提交
111
                global_step += 1
L
lifuchen 已提交
112

L
lifuchen 已提交
113
                #Forward
L
lifuchen 已提交
114 115 116 117 118
                result = model(
                    character,
                    pos_text,
                    mel_pos=pos_mel,
                    length_target=alignment)
L
lifuchen 已提交
119 120 121
                mel_output, mel_output_postnet, duration_predictor_output, _, _ = result
                mel_loss = layers.mse_loss(mel_output, mel)
                mel_postnet_loss = layers.mse_loss(mel_output_postnet, mel)
L
lifuchen 已提交
122 123 124 125
                duration_loss = layers.mean(
                    layers.abs(
                        layers.elementwise_sub(duration_predictor_output,
                                               alignment)))
L
lifuchen 已提交
126 127
                total_loss = mel_loss + mel_postnet_loss + duration_loss

L
lifuchen 已提交
128 129 130 131 132 133 134 135 136 137
                if local_rank == 0:
                    writer.add_scalar('mel_loss',
                                      mel_loss.numpy(), global_step)
                    writer.add_scalar('post_mel_loss',
                                      mel_postnet_loss.numpy(), global_step)
                    writer.add_scalar('duration_loss',
                                      duration_loss.numpy(), global_step)
                    writer.add_scalar('learning_rate',
                                      optimizer._learning_rate.step().numpy(),
                                      global_step)
L
lifuchen 已提交
138

L
lifuchen 已提交
139
                if args.use_data_parallel:
L
lifuchen 已提交
140 141 142 143 144
                    total_loss = model.scale_loss(total_loss)
                    total_loss.backward()
                    model.apply_collective_grads()
                else:
                    total_loss.backward()
L
lifuchen 已提交
145 146 147 148
                optimizer.minimize(
                    total_loss,
                    grad_clip=fluid.dygraph_grad_clip.GradClipByGlobalNorm(cfg[
                        'grad_clip_thresh']))
L
lifuchen 已提交
149 150
                model.clear_gradients()

L
lifuchen 已提交
151 152
                # save checkpoint
                if local_rank == 0 and global_step % args.save_step == 0:
L
lifuchen 已提交
153 154
                    if not os.path.exists(args.save_path):
                        os.mkdir(args.save_path)
L
lifuchen 已提交
155 156
                    save_path = os.path.join(args.save_path,
                                             'fastspeech/%d' % global_step)
L
lifuchen 已提交
157 158
                    dg.save_dygraph(model.state_dict(), save_path)
                    dg.save_dygraph(optimizer.state_dict(), save_path)
L
lifuchen 已提交
159
        if local_rank == 0:
L
lifuchen 已提交
160 161 162
            writer.close()


L
lifuchen 已提交
163
if __name__ == '__main__':
L
lifuchen 已提交
164
    parser = argparse.ArgumentParser(description="Train Fastspeech model")
L
lifuchen 已提交
165
    add_config_options_to_parser(parser)
L
lifuchen 已提交
166 167 168 169
    args = parser.parse_args()
    # Print the whole config setting.
    pprint(args)
    main(args)