forked from KinWaiCheuk/Jointist
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pred_transcription.py
171 lines (141 loc) · 7.05 KB
/
pred_transcription.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
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
from functools import partial
import os
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from pytorch_lightning.plugins import DDPPlugin
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
import End2End.Data as Data
from End2End.tasks.transcription import Transcription, BaselineTranscription
from End2End.models.transcription.seg_baseline import Semantic_Segmentation
from End2End.MIDI_program_map import (
MIDI_Class_NUM,
MIDIClassName2class_idx,
class_idx2MIDIClass,
)
from End2End.data.augmentors import Augmentor
from End2End.lr_schedulers import get_lr_lambda
import End2End.models.transcription.combined as TranscriptionModel
from End2End.losses import get_loss_function
# Libraries related to hydra
import hydra
from hydra.utils import to_absolute_path
@hydra.main(config_path="End2End/config/", config_name="pred_transcription_config")
def main(cfg):
r"""Train an instrument classification system, evluate, and save checkpoints.
Args:
workspace: str, path
config_yaml: str, path
gpus: int
mini_data: bool
Returns:
None
"""
if cfg.MIDI_MAPPING.type=='plugin_names':
cfg.MIDI_MAPPING.plugin_labels_num = PLUGIN_LABELS_NUM
cfg.MIDI_MAPPING.NAME_TO_IX = PLUGIN_LB_TO_IX
cfg.MIDI_MAPPING.IX_TO_NAME = PLUGIN_IX_TO_LB
elif cfg.MIDI_MAPPING.type=='MIDI_class':
cfg.MIDI_MAPPING.plugin_labels_num = MIDI_Class_NUM
cfg.MIDI_MAPPING.NAME_TO_IX = MIDIClassName2class_idx
cfg.MIDI_MAPPING.IX_TO_NAME = class_idx2MIDIClass
else:
raise ValueError(f"Please choose the correct MIDI_MAPPING.type")
# data module
if cfg.datamodule.type=='MSD':
from samidata_pt.loaders.msd.loader import get_dataloader
pred_loader = get_dataloader(**cfg.datamodule.dataloader_cfg)
elif cfg.datamodule.type=='H5Dataset':
h5_root = to_absolute_path(cfg.h5_root)
cfg.datamodule.args.h5_path = os.path.join(h5_root, cfg.h5_name)
dataset = getattr(Data, cfg.datamodule.type)\
(**cfg.datamodule.args, MIDI_MAPPING=cfg.MIDI_MAPPING)
pred_loader = DataLoader(dataset, **cfg.datamodule.dataloader_cfg.pred)
else:
dataset = getattr(Data, cfg.datamodule.type)\
(**cfg.datamodule.args, MIDI_MAPPING=cfg.MIDI_MAPPING)
pred_loader = DataLoader(dataset, **cfg.datamodule.dataloader_cfg.pred)
# model
if cfg.transcription.model.type=='Semantic_Segmentation':
model = Semantic_Segmentation(cfg, **cfg.transcription.model.args)
experiment_name = (
f"Eval-"
f"{cfg.transcription.model.type}-"
f"{cfg.MIDI_MAPPING.type}-"
f"csize={MIDI_Class_NUM}-"
f"bz={cfg.batch_size}"
)
elif cfg.datamodule.type=='H5Dataset':
Model = getattr(TranscriptionModel, cfg.transcription.model.type)
model = Model(cfg, **cfg.transcription.model.args)
h5_filename = os.path.basename(cfg.datamodule.args.h5_path)
experiment_name = (
f"Eval-"
f"{cfg.h5_name}-"
f"{cfg.datamodule.type}-"
f"{cfg.transcription.model.type}-"
f"{cfg.transcription.backend.acoustic.type}-"
f"{cfg.transcription.backend.language.type}_"
f"{cfg.transcription.backend.language.args.hidden_size}-"
f"{cfg.MIDI_MAPPING.type}-"
f"{cfg.inst_sampler.mode}_{cfg.inst_sampler.temp}_{cfg.inst_sampler.samples}inst_"
f"noise{cfg.inst_sampler.audio_noise}-"
f"fps={cfg.transcription.model.args.frames_per_second}-"
f"csize={MIDI_Class_NUM}-"
f"bz={cfg.batch_size}"
)
# loss function
loss_function = get_loss_function(cfg.transcription.model.loss_types)
else:
Model = getattr(TranscriptionModel, cfg.transcription.model.type)
model = Model(cfg, **cfg.transcription.model.args)
experiment_name = (
f"Eval-"
f"{cfg.datamodule.type}-"
f"{cfg.transcription.model.type}-"
f"{cfg.transcription.backend.acoustic.type}-"
f"{cfg.transcription.backend.language.type}_"
f"{cfg.transcription.backend.language.args.hidden_size}-"
f"{cfg.MIDI_MAPPING.type}-"
f"{cfg.inst_sampler.mode}_{cfg.inst_sampler.temp}_{cfg.inst_sampler.samples}inst_"
f"noise{cfg.inst_sampler.audio_noise}-"
f"fps={cfg.transcription.model.args.frames_per_second}-"
f"csize={MIDI_Class_NUM}-"
f"bz={cfg.batch_size}"
)
# loss function
loss_function = get_loss_function(cfg.transcription.model.loss_types)
# callbacks
# save checkpoint callback
lr_monitor = LearningRateMonitor(logging_interval='epoch')
checkpoint_callback = ModelCheckpoint(**cfg.checkpoint,
auto_insert_metric_name=False)
callbacks = [checkpoint_callback, lr_monitor]
logger = pl.loggers.TensorBoardLogger(save_dir='.', name=experiment_name)
checkpoint_path = to_absolute_path(cfg.transcription.evaluation.checkpoint_path)
if cfg.transcription.model.type=='Semantic_Segmentation':
# PL model
pl_model = BaselineTranscription.load_from_checkpoint(checkpoint_path,
network=model,
lr_lambda=None,
batch_data_preprocessor=None,
cfg=cfg)
else:
# PL model
pl_model = Transcription.load_from_checkpoint(checkpoint_path,
network=model,
loss_function=None,
lr_lambda=None,
batch_data_preprocessor=None,
cfg=cfg)
if cfg.trainer.gpus==0: # If CPU is used, disable syncbatch norm
cfg.trainer.sync_batchnorm=False
trainer = pl.Trainer(
**cfg.trainer,
callbacks=callbacks,
plugins=None,
logger=logger
)
# Fit, evaluate, and save checkpoints.
trainer.predict(pl_model, pred_loader)
if __name__ == '__main__':
main()