-
Notifications
You must be signed in to change notification settings - Fork 67
/
train.py
64 lines (49 loc) · 2.06 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
import argparse
import os
import json5
import numpy as np
import torch
from torch.utils.data import DataLoader
from util.utils import initialize_config
def main(config, resume):
torch.manual_seed(config["seed"]) # for both CPU and GPU
np.random.seed(config["seed"])
train_dataloader = DataLoader(
dataset=initialize_config(config["train_dataset"]),
batch_size=config["train_dataloader"]["batch_size"],
num_workers=config["train_dataloader"]["num_workers"],
shuffle=config["train_dataloader"]["shuffle"],
pin_memory=config["train_dataloader"]["pin_memory"]
)
valid_dataloader = DataLoader(
dataset=initialize_config(config["validation_dataset"]),
num_workers=1,
batch_size=1
)
model = initialize_config(config["model"])
optimizer = torch.optim.Adam(
params=model.parameters(),
lr=config["optimizer"]["lr"],
betas=(config["optimizer"]["beta1"], config["optimizer"]["beta2"])
)
loss_function = initialize_config(config["loss_function"])
trainer_class = initialize_config(config["trainer"], pass_args=False)
trainer = trainer_class(
config=config,
resume=resume,
model=model,
loss_function=loss_function,
optimizer=optimizer,
train_dataloader=train_dataloader,
validation_dataloader=valid_dataloader
)
trainer.train()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Wave-U-Net for Speech Enhancement")
parser.add_argument("-C", "--configuration", required=True, type=str, help="Configuration (*.json).")
parser.add_argument("-R", "--resume", action="store_true", help="Resume experiment from latest checkpoint.")
args = parser.parse_args()
configuration = json5.load(open(args.configuration))
configuration["experiment_name"], _ = os.path.splitext(os.path.basename(args.configuration))
configuration["config_path"] = args.configuration
main(configuration, resume=args.resume)