-
Notifications
You must be signed in to change notification settings - Fork 31
/
pit_train.py
43 lines (35 loc) · 1.49 KB
/
pit_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
import sys,os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import argparse
import torch
import torch.multiprocessing as mp
from omegaconf import OmegaConf
from pitch_extend.train import train
torch.backends.cudnn.benchmark = True
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', type=str, required=True,
help="yaml file for configuration")
parser.add_argument('-p', '--checkpoint_path', type=str, default=None,
help="path of checkpoint pt file to resume training")
parser.add_argument('-n', '--name', type=str, required=True,
help="name of the model for logging, saving checkpoint")
args = parser.parse_args()
hp = OmegaConf.load(args.config)
with open(args.config, 'r') as f:
hp_str = ''.join(f.readlines())
assert hp.data.hop_length == 320, \
'hp.data.hop_length must be equal to 320, got %d' % hp.data.hop_length
args.num_gpus = 0
torch.manual_seed(hp.train.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(hp.train.seed)
args.num_gpus = torch.cuda.device_count()
print('Batch size per GPU :', hp.train.batch_size)
if args.num_gpus > 1:
mp.spawn(train, nprocs=args.num_gpus,
args=(args, args.checkpoint_path, hp, hp_str,))
else:
train(0, args, args.checkpoint_path, hp, hp_str)
else:
print('No GPU find!')