-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
328 lines (273 loc) · 11.4 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
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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
import hydra
import os
import shutil
import torch
import warnings
import yaml
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pytorch_lightning as pl
from omegaconf import DictConfig, OmegaConf, open_dict
from tqdm import tqdm
import alignnet
# Load clearml Task only if clearml is imported
try:
from clearml import Task
from clearml.backend_api.session.defs import MissingConfigError
except ModuleNotFoundError as err:
def Task(**kwargs):
return None
def post_train(model, audio_data, loggers, task=None):
audio_data.batch_size = 1
data_loaders = {
"train": audio_data.train_dataloader(),
"val": audio_data.val_dataloader(),
"test": audio_data.test_dataloader(),
}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()
# Save estimations
out_df = pd.DataFrame()
with torch.no_grad():
for dataset_split, loader in data_loaders.items():
print(f"Loading dataset: {dataset_split}")
results = []
for audio, mos, dataset in tqdm(loader):
audio = audio.to(device)
dataset = dataset.to(device)
est = model(audio, dataset)
est = est.numpy(force=True)
est = np.squeeze(est)
audio_net_est = model.network.audio_net(audio)
audio_net_est = audio_net_est.numpy(force=True)
audio_net_est = np.squeeze(audio_net_est)
mos = mos.numpy(force=True)
mos = np.squeeze(mos)
results.append(
[float(mos), float(est), int(dataset), float(audio_net_est)]
)
results = np.array(results)
data_df = pd.DataFrame(
results,
columns=["MOS", "Estimation", "Dataset_Index", "AudioNet_Estimation"],
)
data_df["Dataset"] = dataset_split
out_df = pd.concat([out_df, data_df])
corr = np.corrcoef(results[:, 0], results[:, 1])[0, 1]
print(f"{dataset_split} corr coef: {corr:6f}")
metric_name = f"test_pearsons/{dataset_split}"
metrics = dict()
metrics[metric_name] = corr
rmse = np.sqrt(np.mean((results[:, 0] - results[:, 1]) ** 2))
metrics[f"test_rmse/{dataset_split}"] = rmse
if dataset_split == "test":
for dix in np.unique(data_df["Dataset_Index"]):
df_sub = data_df[data_df["Dataset_Index"] == dix]
corr = np.corrcoef(df_sub["MOS"], df_sub["Estimation"])[0, 1]
rmse = np.sqrt(np.mean((df_sub["MOS"] - df_sub["Estimation"]) ** 2))
corr_name = f"test_pearsons/dataset {dix}"
metrics[corr_name] = corr
rmse_name = f"test_rmse/dataset {dix}"
metrics[rmse_name] = rmse
[logger.log_metrics(metrics) for logger in loggers]
estimations_name = "estimations.csv"
out_df.to_csv(estimations_name, index=False)
# Store estimations to clearml
if task is not None:
task.upload_artifact(artifact_object=estimations_name, name="estimations csv")
task.upload_artifact(artifact_object=out_df, name="estimations df")
colormap = cm.rainbow
# Plot estimations vs MOS
for k, ds in enumerate(np.unique(out_df["Dataset"])):
df_sub = out_df[out_df["Dataset"] == ds]
mos = df_sub["MOS"]
est = df_sub["Estimation"]
dataset_index = df_sub["Dataset_Index"]
plt.plot([1, 5], [1, 5], color="black", linestyle="dashed")
plt.scatter(x=mos, y=est, c=dataset_index, alpha=0.1, cmap=colormap)
corrcoef = np.corrcoef(mos, est)[0, 1]
rmse = np.sqrt(np.mean((mos - est) ** 2))
title_str = f"{ds} set, LCC={corrcoef:.4f}, RMSE={rmse:.4f}"
for dx in np.unique(dataset_index):
dx_ix = dataset_index == dx
mos_dx = mos[dx_ix]
est_dx = est[dx_ix]
corrcoef_dx = np.corrcoef(mos_dx, est_dx)[0, 1]
rmse_dx = np.sqrt(np.mean((mos_dx - est_dx) ** 2))
subtitle_str = (
f", (Dataset {dx}, LCC={corrcoef_dx:.4f}, RMSE={rmse_dx:.4f})"
)
title_str += subtitle_str
plt.title(title_str)
plt.xlabel("MOS")
plt.ylabel("Estimation")
plt.show()
# audio_net vs aligner estimations
for k, ds in enumerate(np.unique(out_df["Dataset"])):
df_sub = out_df[out_df["Dataset"] == ds]
mos = df_sub["AudioNet_Estimation"]
est = df_sub["Estimation"]
dataset_index = df_sub["Dataset_Index"]
plt.plot([1, 5], [1, 5], color="black", linestyle="dashed")
plt.scatter(x=mos, y=est, c=dataset_index, alpha=0.1, cmap=colormap)
plt.title(f"{ds} set audio_net vs estimation ")
plt.xlabel("audio_net estimation")
plt.ylabel("final estimation")
plt.show()
for k, dx in enumerate(np.unique(out_df["Dataset_Index"])):
xv = np.arange(0.5, 5.5, step=0.01)
xv = torch.Tensor(xv)
xv = xv[:, None]
xv = xv.to(device)
data_tensor = dx * torch.ones(xv.shape)
data_tensor = data_tensor.squeeze()
data_tensor = data_tensor.to(int)
data_tensor = data_tensor.to(device)
yv = model.network.aligner(xv, data_tensor)
xv = xv.cpu().detach().numpy()
yv = yv.cpu().detach().numpy()
plt.plot([1, 5], [1, 5], color="black", linestyle="dashed")
plt.scatter(xv, yv)
plt.title(f"Alignment function for dataset {dx}")
plt.xlabel("Raw score")
plt.ylabel("Aligned")
plt.show()
@hydra.main(config_path="alignnet/config", config_name="conf.yaml", version_base=None)
def main(cfg: DictConfig) -> None:
if cfg.logging.logger == "clearml":
try:
task = Task.init(
project_name=cfg.project.name,
task_name=cfg.project.task,
)
except MissingConfigError as E:
print(f"{E}")
print(
f"If you do not want to install clearML and want to avoid this error in the future, set `logging=none` override."
)
task = None
else:
task = None
print("Working directory : {}".format(os.getcwd()))
# Seed
seed = cfg.common.seed
if seed is None:
rng = np.random.default_rng()
seed = rng.choice(10000)
cfg.common.seed = seed
pl.seed_everything(seed)
# Transform
transform = hydra.utils.instantiate(cfg.transform)
data_class = hydra.utils.instantiate(cfg.dataclass)
audio_data = hydra.utils.instantiate(
cfg.data, transform=transform, DataClass=data_class
)
num_datasets = len(cfg.data.data_dirs)
# Lightning logs
# Initialize tensorboard logger, letting hydra control directory and versions
tb_logger = pl.loggers.TensorBoardLogger(
save_dir=".",
name="",
version="",
)
loggers = [tb_logger]
[logger.log_hyperparams(dict(cfg)) for logger in loggers]
checkpoint_callback = hydra.utils.instantiate(cfg.checkpoint)
callbacks = [checkpoint_callback]
if "earlystop" in cfg:
# Earlystop needs monitor (e.g., val-loss) and mode (e.g., min). This can be added via CLI/cfg. Otherwise steal the checkpoint values.
stop_params = {"monitor": None, "mode": None}
for k, _ in stop_params.items():
if k in cfg.earlystop:
stop_params[k] = cfg.earlystop[k]
else:
stop_params[k] = cfg.checkpoint[k]
early_stopping_callback = hydra.utils.instantiate(cfg.earlystop, **stop_params)
callbacks.append(early_stopping_callback)
# Trainer
trainer = hydra.utils.instantiate(
cfg.optimization, callbacks=callbacks, logger=loggers
)
num_datasets = len(cfg.data.data_dirs)
# Initialize network
network = hydra.utils.instantiate(
cfg.network, aligner={"num_datasets": num_datasets}
)
loss = hydra.utils.instantiate(cfg.loss)
optimizer = hydra.utils.instantiate(cfg.optimizer, lr=cfg.common.lr)
# initialize model
if cfg.finetune.restore_file is not None:
print(f"Loading model from checkpoint: {cfg.finetune.restore_file}")
# initialize model
model_class = hydra.utils.get_class(cfg.model._target_)
# Path to pretrained model checkpoint
model_path = os.path.join(cfg.finetune.restore_file, "model.ckpt")
restore_cfg_path = os.path.join(cfg.finetune.restore_file, "config.yaml")
with open(restore_cfg_path, "r") as f:
restore_yaml = yaml.safe_load(f)
restore_cfg = DictConfig(restore_yaml)
restore_network = hydra.utils.instantiate(restore_cfg.network)
# Turn restored audio_net gradients on or off depending on new network settings
old_freeze_name = "audio_net_freeze_steps"
if hasattr(network, old_freeze_name):
frozen_steps = getattr(network, old_freeze_name)
else:
frozen_steps = network.audio_net_freeze_epochs
if frozen_steps > 0:
restore_network.set_audio_net_update_status(False)
else:
restore_network.set_audio_net_update_status(True)
# Initialize identical network to pretrained version (necessary to appropriately load in aligner)
# aligner is not transferable (different sizes based on number of datasets)
restored_model = model_class.load_from_checkpoint(
model_path, network=restore_network, loss=loss, optimizer=optimizer
)
# Grab audio_net from checkpoint
network.audio_net = restored_model.network.audio_net
model = hydra.utils.instantiate(
cfg.model, network=network, loss=loss, optimizer=optimizer
)
print(model)
# Add working directory to config
with open_dict(cfg):
cfg.project.working_dir = os.getcwd()
# Save a version of the config
cfg_yaml = OmegaConf.to_yaml(cfg)
cfg_out = "input_config.yaml"
with open(cfg_out, "w") as file:
file.write(cfg_yaml)
if cfg.common.auto_batch_size:
tuner = pl.tuner.Tuner(trainer)
tuner.scale_batch_size(model, datamodule=audio_data)
# Fit Trainer
trainer.fit(model, audio_data)
best_model_path = trainer.checkpoint_callback.best_model_path
trained_model_path = "trained_model"
os.makedirs(trained_model_path)
# Save another copy of the top model
top_model_path = os.path.join(trained_model_path, "model.ckpt")
shutil.copy(best_model_path, top_model_path)
print(f'experiment_path = "{os.getcwd()}"')
print(f'model_ckpt = "{best_model_path}"')
# Create output config
output_config = DictConfig({})
output_config.model = cfg.model
output_config.network = cfg.network
# Store num datasets directly
output_config.network.aligner.num_datasets = num_datasets
# Convert to yaml
output_config = OmegaConf.to_yaml(output_config)
# Save output config
output_config_path = os.path.join(trained_model_path, "config.yaml")
with open(output_config_path, "w") as file:
file.write(output_config)
# Get model class
model_class = hydra.utils.get_class(cfg.model._target_)
# Load best model
model = model_class.load_from_checkpoint(best_model_path, network=network)
post_train(model, audio_data, loggers, task=task)
if __name__ == "__main__":
main()