-
Notifications
You must be signed in to change notification settings - Fork 19
/
train.py
80 lines (68 loc) · 2.7 KB
/
train.py
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
import os
import argparse
import torch
try:
import torch_musa
except ImportError:
pass
from torch.optim import lr_scheduler
from logger import utils
from reflow.data_loaders import get_data_loaders
from reflow.vocoder import Vocoder, Unit2Wav_VAE
from reflow.solver import train
def parse_args(args=None, namespace=None):
"""Parse command-line arguments."""
parser = argparse.ArgumentParser()
parser.add_argument(
"-c",
"--config",
type=str,
required=True,
help="path to the config file")
return parser.parse_args(args=args, namespace=namespace)
if __name__ == '__main__':
# parse commands
cmd = parse_args()
# load config
args = utils.load_config(cmd.config)
print(' > config:', cmd.config)
print(' > exp:', args.env.expdir)
# load vocoder
vocoder = Vocoder(args.vocoder.type, args.vocoder.ckpt, device=args.device)
# load model
if args.model.type == 'RectifiedFlow_VAE':
model = Unit2Wav_VAE(
args.data.sampling_rate,
args.data.block_size,
args.model.win_length,
args.data.encoder_out_channels,
args.model.n_spk,
args.model.use_pitch_aug,
vocoder.dimension,
args.model.n_layers,
args.model.n_chans,
args.model.n_hidden,
args.model.back_bone,
args.model.use_attention)
else:
raise ValueError(f" [x] Unknown Model: {args.model.type}")
# load parameters
optimizer = torch.optim.AdamW(model.parameters())
initial_global_step, model, optimizer = utils.load_model(args.env.expdir, model, optimizer, device=args.device)
for param_group in optimizer.param_groups:
param_group['initial_lr'] = args.train.lr
param_group['lr'] = args.train.lr * args.train.gamma ** max((initial_global_step - 2) // args.train.decay_step, 0)
param_group['weight_decay'] = args.train.weight_decay
scheduler = lr_scheduler.StepLR(optimizer, step_size=args.train.decay_step, gamma=args.train.gamma, last_epoch=initial_global_step-2)
# device
if args.device == 'cuda':
torch.cuda.set_device(args.env.gpu_id)
model.to(args.device)
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.to(args.device)
# datas
loader_train, loader_valid = get_data_loaders(args, whole_audio=False)
# run
train(args, initial_global_step, model, optimizer, scheduler, vocoder, loader_train, loader_valid)