def main(config, resume):
torch.manual_seed(config["seed"]) # both CPU and CUDA
np.random.seed(config["seed"])
test_dataloader = DataLoader(
dataset=initialize_config(config["test_dataset"]),
batch_size=1,
num_workers=1,
)
generator = initialize_config(config["generator_model"])
discriminator = initialize_config(config["discriminator_model"])
generator_optimizer = torch.optim.Adam(
params=generator.parameters(),
lr=config["optimizer"]["G_lr"],
betas=(config["optimizer"]["beta1"], config["optimizer"]["beta2"]))
discriminator_optimizer = torch.optim.Adam(
params=discriminator.parameters(),
lr=config["optimizer"]["D_lr"],
betas=(config["optimizer"]["beta1"], config["optimizer"]["beta2"]))
additional_loss_function = initialize_config(
config["additional_loss_function"])
tester = Tester(config=config,
resume=resume,
generator=generator,
discriminator=discriminator,
generator_optimizer=generator_optimizer,
discriminator_optimizer=discriminator_optimizer,
additional_loss_function=additional_loss_function,
test_dl=test_dataloader)
tester.test()
def main(config, checkpoint_path, output_dir):
inferencer_class = initialize_config(config["inference"], pass_args=False)
inferencer = inferencer_class(
config,
checkpoint_path,
output_dir
)
inferencer.inference()
def main(config, resume):
torch.manual_seed(config["seed"]) # both CPU and CUDA
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"] # Very small data set False
)
validation_dataloader = DataLoader(
dataset=initialize_config(config["validation_dataset"]),
batch_size=1,
num_workers=1,
)
generator = initialize_config(config["generator_model"])
discriminator = initialize_config(config["discriminator_model"])
generator_optimizer = torch.optim.Adam(
params=generator.parameters(),
lr=config["optimizer"]["G_lr"],
betas=(config["optimizer"]["beta1"], config["optimizer"]["beta2"]))
discriminator_optimizer = torch.optim.Adam(
params=discriminator.parameters(),
lr=config["optimizer"]["D_lr"],
betas=(config["optimizer"]["beta1"], config["optimizer"]["beta2"]))
additional_loss_function = initialize_config(
config["additional_loss_function"])
trainer = Trainer(
config=config,
resume=resume,
generator=generator,
discriminator=discriminator,
generator_optimizer=generator_optimizer,
discriminator_optimizer=discriminator_optimizer,
additional_loss_function=additional_loss_function,
train_dl=train_dataloader,
validation_dl=validation_dataloader,
)
trainer.train()
def _load_dataloader(dataset_config):
dataset = initialize_config(dataset_config)
dataloader = DataLoader(
dataset=dataset,
batch_size=1,
num_workers=0,
)
return dataloader
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()
def main(config, resume):
torch.manual_seed(config["seed"]) # 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"], # Set it to False for very small dataset, otherwise True.
collate_fn=pad_to_longest_in_one_batch)
validation_dataloader = DataLoader(
dataset=initialize_config(config["validation_dataset"]),
batch_size=1,
num_workers=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 = Trainer(
config=config,
resume=resume,
model=model,
optimizer=optimizer,
loss_function=loss_function,
train_dl=train_dataloader,
validation_dl=validation_dataloader,
)
trainer.train()
def main(config, resume):
torch.manual_seed(int(config["seed"])) # both CPU and CUDA
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"] # Very small data set False
)
validation_dataloader = DataLoader(dataset=initialize_config(
config["validation_dataset"]),
batch_size=1,
num_workers=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 = GeneralTrainer(config=config,
resume=resume,
model=model,
optim=optimizer,
loss_fucntion=loss_function,
train_dl=train_dataloader,
validation_dl=validation_dataloader)
trainer.train()
def _load_model(model_config, checkpoint_path, device):
model = initialize_config(model_config)
if os.path.splitext(os.path.basename(checkpoint_path))[-1] == ".tar":
model_checkpoint = torch.load(checkpoint_path, map_location=device)
model_static_dict = model_checkpoint["model"]
print(
f"Loading model checkpoint with *.tar format, the epoch is: {model_checkpoint['epoch']}."
)
else:
model_static_dict = torch.load(checkpoint_path,
map_location=device)
model.load_state_dict(model_static_dict)
model.to(device)
model.eval()
return model
args = parser.parse_args()
"""
Preparation
"""
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
config = json.load(open(args.config))
model_checkpoint_path = args.model_checkpoint_path
output_dir = args.output_dir
assert os.path.exists(output_dir), "Enhanced directory should be exist."
"""
DataLoader
"""
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
dataloader = DataLoader(dataset=initialize_config(config["dataset"]), batch_size=1, num_workers=0)
"""
Model
"""
model = initialize_config(config["model"])
model.load_state_dict(load_checkpoint(model_checkpoint_path, device))
model.to(device)
model.eval()
"""
Enhancement
"""
sample_length = config["custom"]["sample_length"]
for mixture, name in tqdm(dataloader):
assert len(name) == 1, "Only support batch size is 1 in enhancement stage."
help="Checkpoint.")
args = parser.parse_args()
"""
Preparation
"""
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
config = json.load(open(args.config))
model_checkpoint_path = args.model_checkpoint_path
output_dir = args.output_dir
assert os.path.exists(output_dir), "Enhanced directory should be exist."
"""
DataLoader
"""
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device(
"cpu")
dataloader = DataLoader(dataset=initialize_config(config["dataset"]),
batch_size=1,
num_workers=0)
"""
Model
"""
model = initialize_config(config["model"])
model.load_state_dict(load_checkpoint(model_checkpoint_path, device))
model.to(device)
model.eval()
"""
Enhancement
"""
sample_length = dataloader.dataset.sample_length
for mixture, name in tqdm(dataloader):
assert len(name) == 1, "Only support batch size is 1 in enhancement stage."
def main(config, resume):
# Random seed for both CPU and GPU.
torch.manual_seed(config["seed"])
np.random.seed(config["seed"])
def collate_fn_pad(batch):
"""
Returns:
[B, F, T (Longest)]
"""
noisy_list = []
clean_list = []
n_frames_list = []
names = []
for noisy, clean, n_frames, name in batch:
noisy_list.append(torch.tensor(noisy).permute(
1, 0)) # [F, T] => [T, F]
clean_list.append(torch.tensor(clean).permute(
1, 0)) # [1, T] => [T, 1]
n_frames_list.append(n_frames)
names.append(name)
# seq_list = [(T1, F), (T2, F), ...]
# item.size() must be (T, *)
# return (longest_T, len(seq_list), *)
noisy_list = pad_sequence(noisy_list).permute(
1, 2, 0) # ([T1, F], [T2, F], ...) => [T, B, F] => [B, F, T]
clean_list = pad_sequence(clean_list).permute(
1, 2, 0) # ([T1, 1], [T2, 1], ...) => [T, B, 1] => [B, 1, T]
return noisy_list, clean_list, n_frames_list, names
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"],
collate_fn=collate_fn_pad)
valid_dataloader = DataLoader(dataset=initialize_config(
config["validation_dataset"]),
num_workers=0,
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()
