LC-beating/train.py at infer · meteroad/LC-beating · GitHub

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
import argparse
from trainer import Trainer

from param_file.params import *
from simple_test import test_func

def train(args):
    # train or test
    if args.mode == 'train':
        ## -- create log file -- ##
        log_file_name = 'log/' + args.network + '_' + args.mark + '.log'
        args.log_file = open(log_file_name, 'ab') ##
        # arg_settings_file =

        trainer = Trainer(args)
        trainer.train()
    elif args.mode == 'test':
        tester = test_func(args)
        tester.test()
    elif args.mode == "real_time":
        tester = test_func(args)
        tester.calculate_rtf()
    return 0

if __name__ == '__main__':
    # torch.multiprocessing.set_start_method('spawn')

    parser = argparse.ArgumentParser()

    # training level args input
    parser.add_argument('-n', "--network", default='lc42_dt_1')
    parser.add_argument('-m', "--mark", default='round4_device1')

    parser.add_argument('-md', "--mode", default='train') # real_time, train, test

    parser.add_argument('--eval_mode', default = "online") # online, offline
    parser.add_argument('--inference_model', default = "simple") # simple, DBN， dl, PF
    parser.add_argument('--peak_type', default = "simple") # librosa, simple, thresh
    parser.add_argument('--peak_latency', type=int, default=5)

    # testing level args settings
    parser.add_argument("--check_num", default = "135") # 0 refer to the latest file

    # eight fold validation num
    parser.add_argument("--folder", type=int, default = 0, help="-1,0,1,2,3,4,5,6,7")

    # training/testing mode
    parser.add_argument("--only_beat", default = False)
    parser.add_argument("--only_downbeat", default = False)##

    # specific training settings
    parser.add_argument("-me","--max_epochs", type=int, default=250)
    parser.add_argument("--patience", type=int, default=50)

    parser.add_argument('--train_subset', type=str, default='train')
    parser.add_argument('--val_subset', type=str, default='val')

    parser.add_argument('--batch_size', type=int, default=10)
    parser.add_argument('--lr', type=float, default=0.001)
    parser.add_argument('--device', type=str, default="cuda")
    parser.add_argument('--device_index', type=int, default=1)
    parser.add_argument('--num_workers', type=int, default=4) # 4

    # loss settings
    parser.add_argument('--loss', type=str, default="softbce_n")
    parser.add_argument('--loss_param', type=float, default=10)
    parser.add_argument("--clip_grad", type=float, default=0.5) # no use

    parser.add_argument('--checkname', type=int, default=210)
    parser.add_argument('--ifcheckpoint', default=True)

    # test mode, small training epoch
    parser.add_argument("--epoch_batch", type=int, default=20)
    parser.add_argument('--iftestmode', default=False)


    args = parser.parse_args()

    args.length_sec = length_sec
    args.train_length = train_length
    args.eval_length = eval_length
    args.target_length = target_frame
    args.limit_length = limit_length

    args.train_sets = train_sets
    args.test_sets = test_sets
    args.model_path = MODEL_PATH
    args.pp_model_path = MODEL_PATH_PP
    args.scalar_writer = SCALAR_WRITER
    args.scalar_writer_pp = SCALAR_WRITER_PP

    args.sr = target_sr
    args.hop_length = hop_length
    args.target_factor = hop_length

    args.root_path = root_path

    train(args)