def test_upload():
# Make dirs
os.makedirs("tmp/publish_dir", exist_ok=True)
populate_wham_dir("tmp/wham")
# Dataset and NN
train_set = WhamDataset("tmp/wham", task="sep_clean")
model = ConvTasNet(n_src=2,
n_repeats=2,
n_blocks=2,
bn_chan=16,
hid_chan=4,
skip_chan=8,
n_filters=32)
# Save publishable
model_conf = model.serialize()
model_conf.update(train_set.get_infos())
save_publishable("tmp/publish_dir", model_conf, metrics={}, train_conf={})
# Upload
token = os.getenv("ACCESS_TOKEN")
if token: # ACESS_TOKEN is not available on forks.
zen, current = upload_publishable(
"tmp/publish_dir",
uploader="Manuel Pariente",
affiliation="INRIA",
use_sandbox=True,
unit_test=True, # Remove this argument and monkeypatch `input()`
git_username="******",
)
# Assert metadata is correct
meta = current.json()["metadata"]
assert meta["creators"][0]["name"] == "Manuel Pariente"
assert meta["creators"][0]["affiliation"] == "INRIA"
assert "asteroid-models" in [
d["identifier"] for d in meta["communities"]
]
# Clean up
zen.remove_deposition(current.json()["id"])
shutil.rmtree("tmp/wham")
def test_upload():
# Make dirs
os.makedirs('tmp/publish_dir', exist_ok=True)
populate_wham_dir('tmp/wham')
# Dataset and NN
train_set = WhamDataset('tmp/wham', task='sep_clean')
model = ConvTasNet(n_src=2,
n_repeats=2,
n_blocks=2,
bn_chan=16,
hid_chan=4,
skip_chan=8,
n_filters=32)
# Save publishable
model_conf = model.serialize()
model_conf.update(train_set.get_infos())
save_publishable('tmp/publish_dir', model_conf, metrics={}, train_conf={})
# if False:
# Upload
zen, current = upload_publishable(
'tmp/publish_dir',
uploader="Manuel Pariente",
affiliation="INRIA",
use_sandbox=True,
unit_test=True, # Remove this argument and monkeypatch `input()`
)
# Assert metadata is correct
meta = current.json()['metadata']
assert meta['creators'][0]['name'] == "Manuel Pariente"
assert meta['creators'][0]['affiliation'] == "INRIA"
assert 'asteroid-models' in [d['identifier'] for d in meta['communities']]
# Clean up
zen.remove_deposition(current.json()['id'])
shutil.rmtree('tmp/wham')
def main(conf):
train_set = WhamDataset(
conf["data"]["train_dir"],
conf["data"]["task"],
sample_rate=conf["data"]["sample_rate"],
segment=conf["data"]["segment"],
nondefault_nsrc=conf["data"]["nondefault_nsrc"],
)
val_set = WhamDataset(
conf["data"]["valid_dir"],
conf["data"]["task"],
sample_rate=conf["data"]["sample_rate"],
nondefault_nsrc=conf["data"]["nondefault_nsrc"],
)
train_loader = DataLoader(
train_set,
shuffle=True,
batch_size=conf["training"]["batch_size"],
num_workers=conf["training"]["num_workers"],
drop_last=True,
)
val_loader = DataLoader(
val_set,
shuffle=False,
batch_size=conf["training"]["batch_size"],
num_workers=conf["training"]["num_workers"],
drop_last=True,
)
# Update number of source values (It depends on the task)
conf["masknet"].update({"n_src": train_set.n_src})
model = DPTNet(**conf["filterbank"], **conf["masknet"])
optimizer = make_optimizer(model.parameters(), **conf["optim"])
from asteroid.engine.schedulers import DPTNetScheduler
schedulers = {
"scheduler":
DPTNetScheduler(optimizer,
len(train_loader) // conf["training"]["batch_size"],
64),
"interval":
"step",
}
# Just after instantiating, save the args. Easy loading in the future.
exp_dir = conf["main_args"]["exp_dir"]
os.makedirs(exp_dir, exist_ok=True)
conf_path = os.path.join(exp_dir, "conf.yml")
with open(conf_path, "w") as outfile:
yaml.safe_dump(conf, outfile)
# Define Loss function.
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx")
system = System(
model=model,
loss_func=loss_func,
optimizer=optimizer,
scheduler=schedulers,
train_loader=train_loader,
val_loader=val_loader,
config=conf,
)
# Define callbacks
checkpoint_dir = os.path.join(exp_dir, "checkpoints/")
checkpoint = ModelCheckpoint(checkpoint_dir,
monitor="val_loss",
mode="min",
save_top_k=5,
verbose=True)
early_stopping = False
if conf["training"]["early_stop"]:
early_stopping = EarlyStopping(monitor="val_loss",
patience=30,
verbose=True)
# Don't ask GPU if they are not available.
gpus = -1 if torch.cuda.is_available() else None
trainer = pl.Trainer(
max_epochs=conf["training"]["epochs"],
checkpoint_callback=checkpoint,
early_stop_callback=early_stopping,
default_root_dir=exp_dir,
gpus=gpus,
distributed_backend="ddp",
gradient_clip_val=conf["training"]["gradient_clipping"],
)
trainer.fit(system)
best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
with open(os.path.join(exp_dir, "best_k_models.json"), "w") as f:
json.dump(best_k, f, indent=0)
state_dict = torch.load(checkpoint.best_model_path)
system.load_state_dict(state_dict=state_dict["state_dict"])
system.cpu()
to_save = system.model.serialize()
to_save.update(train_set.get_infos())
torch.save(to_save, os.path.join(exp_dir, "best_model.pth"))
def main(conf):
train_set = WhamDataset(conf['data']['train_dir'],
conf['data']['task'],
sample_rate=conf['data']['sample_rate'],
nondefault_nsrc=conf['data']['nondefault_nsrc'])
val_set = WhamDataset(conf['data']['valid_dir'],
conf['data']['task'],
sample_rate=conf['data']['sample_rate'],
nondefault_nsrc=conf['data']['nondefault_nsrc'])
train_loader = DataLoader(train_set,
shuffle=True,
batch_size=conf['training']['batch_size'],
num_workers=conf['training']['num_workers'],
drop_last=True)
val_loader = DataLoader(val_set,
shuffle=False,
batch_size=conf['training']['batch_size'],
num_workers=conf['training']['num_workers'],
drop_last=True)
# Update number of source values (It depends on the task)
conf['masknet'].update({'n_src': train_set.n_src})
# Define model and optimizer
model = ConvTasNet(**conf['filterbank'], **conf['masknet'])
optimizer = make_optimizer(model.parameters(), **conf['optim'])
# Define scheduler
scheduler = None
if conf['training']['half_lr']:
scheduler = ReduceLROnPlateau(optimizer=optimizer,
factor=0.5,
patience=5)
# Just after instantiating, save the args. Easy loading in the future.
exp_dir = conf['main_args']['exp_dir']
os.makedirs(exp_dir, exist_ok=True)
conf_path = os.path.join(exp_dir, 'conf.yml')
with open(conf_path, 'w') as outfile:
yaml.safe_dump(conf, outfile)
# Define Loss function.
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from='pw_mtx')
system = System(model=model,
loss_func=loss_func,
optimizer=optimizer,
train_loader=train_loader,
val_loader=val_loader,
scheduler=scheduler,
config=conf)
# Define callbacks
checkpoint_dir = os.path.join(exp_dir, 'checkpoints/')
checkpoint = ModelCheckpoint(checkpoint_dir,
monitor='val_loss',
mode='min',
save_top_k=5,
verbose=1)
early_stopping = False
if conf['training']['early_stop']:
early_stopping = EarlyStopping(monitor='val_loss',
patience=10,
verbose=1)
# Don't ask GPU if they are not available.
gpus = -1 if torch.cuda.is_available() else None
trainer = pl.Trainer(
max_epochs=conf['training']['epochs'],
checkpoint_callback=checkpoint,
early_stop_callback=early_stopping,
default_save_path=exp_dir,
gpus=gpus,
distributed_backend='dp',
train_percent_check=1.0, # Useful for fast experiment
gradient_clip_val=5.)
trainer.fit(system)
best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
with open(os.path.join(exp_dir, "best_k_models.json"), "w") as f:
json.dump(best_k, f, indent=0)
# Save best model (next PL version will make this easier)
best_path = [b for b, v in best_k.items() if v == min(best_k.values())][0]
state_dict = torch.load(best_path)
system.load_state_dict(state_dict=state_dict['state_dict'])
system.cpu()
to_save = system.model.serialize()
to_save.update(train_set.get_infos())
torch.save(to_save, os.path.join(exp_dir, 'best_model.pth'))
def main(conf):
assert (conf["training"]["batch_size"] %
2 == 0), "Batch size must be divisible by two to run this recipe"
train_set = WhamDataset(
conf["data"]["train_dir"],
"sep_clean",
sample_rate=conf["data"]["sample_rate"],
segment=conf["data"]["segment"],
nondefault_nsrc=None,
)
val_set = WhamDataset(
conf["data"]["valid_dir"],
"sep_clean",
sample_rate=conf["data"]["sample_rate"],
nondefault_nsrc=None,
)
train_loader = DataLoader(
train_set,
shuffle=True,
batch_size=conf["training"]["batch_size"],
num_workers=conf["training"]["num_workers"],
drop_last=True,
)
val_loader = DataLoader(
val_set,
shuffle=False,
batch_size=conf["training"]["batch_size"],
num_workers=conf["training"]["num_workers"],
drop_last=True,
)
model = DPRNNTasNet(**conf["filterbank"],
**conf["masknet"],
sample_rate=conf["data"]["sample_rate"])
optimizer = make_optimizer(model.parameters(), **conf["optim"])
# Define scheduler
scheduler = None
if conf["training"]["half_lr"]:
scheduler = ReduceLROnPlateau(optimizer=optimizer,
factor=0.5,
patience=5)
# Just after instantiating, save the args. Easy loading in the future.
exp_dir = conf["main_args"]["exp_dir"]
os.makedirs(exp_dir, exist_ok=True)
conf_path = os.path.join(exp_dir, "conf.yml")
with open(conf_path, "w") as outfile:
yaml.safe_dump(conf, outfile)
# Define Loss function.
loss_func = {
"pit": PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx"),
"mixit": MixITLossWrapper(pairwise_neg_sisdr, generalized=True),
}
system = MixITSystem(
model=model,
loss_func=loss_func,
optimizer=optimizer,
train_loader=train_loader,
val_loader=val_loader,
scheduler=scheduler,
config=conf,
)
# Define callbacks
callbacks = []
checkpoint_dir = os.path.join(exp_dir, "checkpoints/")
checkpoint = ModelCheckpoint(checkpoint_dir,
monitor="val_loss",
mode="min",
save_top_k=5,
verbose=True)
callbacks.append(checkpoint)
if conf["training"]["early_stop"]:
callbacks.append(
EarlyStopping(monitor="val_loss",
mode="min",
patience=30,
verbose=True))
# Don't ask GPU if they are not available.
gpus = -1 if torch.cuda.is_available() else None
distributed_backend = "ddp" if torch.cuda.is_available() else None
trainer = pl.Trainer(
max_epochs=conf["training"]["epochs"],
callbacks=callbacks,
default_root_dir=exp_dir,
gpus=gpus,
distributed_backend=distributed_backend,
gradient_clip_val=conf["training"]["gradient_clipping"],
)
trainer.fit(system)
best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
with open(os.path.join(exp_dir, "best_k_models.json"), "w") as f:
json.dump(best_k, f, indent=0)
state_dict = torch.load(checkpoint.best_model_path)
system.load_state_dict(state_dict=state_dict["state_dict"])
system.cpu()
to_save = system.model.serialize()
to_save.update(train_set.get_infos())
torch.save(to_save, os.path.join(exp_dir, "best_model.pth"))
def main(conf):
train_set = WhamDataset(
conf["data"]["train_dir"],
conf["data"]["task"],
sample_rate=conf["data"]["sample_rate"],
segment=conf["data"]["segment"],
nondefault_nsrc=conf["data"]["nondefault_nsrc"],
)
val_set = WhamDataset(
conf["data"]["valid_dir"],
conf["data"]["task"],
sample_rate=conf["data"]["sample_rate"],
nondefault_nsrc=conf["data"]["nondefault_nsrc"],
)
train_loader = DataLoader(
train_set,
shuffle=True,
batch_size=conf["training"]["batch_size"],
num_workers=conf["training"]["num_workers"],
drop_last=True,
)
val_loader = DataLoader(
val_set,
shuffle=False,
batch_size=conf["training"]["batch_size"],
num_workers=conf["training"]["num_workers"],
drop_last=True,
)
# Update number of source values (It depends on the task)
conf["masknet"].update({"n_src": train_set.n_src})
model = DPRNNTasNet(**conf["filterbank"], **conf["masknet"])
optimizer = make_optimizer(model.parameters(), **conf["optim"])
# Define scheduler
scheduler = None
if conf["training"]["half_lr"]:
scheduler = ReduceLROnPlateau(optimizer=optimizer,
factor=0.5,
patience=5)
# Just after instantiating, save the args. Easy loading in the future.
exp_dir = conf["main_args"]["exp_dir"]
os.makedirs(exp_dir, exist_ok=True)
conf_path = os.path.join(exp_dir, "conf.yml")
with open(conf_path, "w") as outfile:
yaml.safe_dump(conf, outfile)
# Define Loss function.
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx")
system = System(
model=model,
loss_func=loss_func,
optimizer=optimizer,
train_loader=train_loader,
val_loader=val_loader,
scheduler=scheduler,
config=conf,
)
# Define callbacks
checkpoint_dir = os.path.join(exp_dir, 'checkpoints/')
checkpoint = ModelCheckpoint(checkpoint_dir,
monitor='val_loss',
verbose=True,
mode='min',
save_top_k=5)
early_stopping = False
if conf["training"]["early_stop"]:
early_stopping = EarlyStopping(monitor="val_loss",
patience=30,
verbose=1)
# Don't ask GPU if they are not available.
gpus = -1 if torch.cuda.is_available() else None
trainer = pl.Trainer(
max_epochs=conf['training']['epochs'],
checkpoint_callback=checkpoint,
early_stop_callback=early_stopping,
default_root_dir=exp_dir,
gpus=gpus,
distributed_backend='ddp',
gradient_clip_val=conf['training']["gradient_clipping"])
trainer.fit(system)
state_dict = torch.load(checkpoint.best_model_path)
system.load_state_dict(state_dict=state_dict['state_dict'])
system.cpu()
to_save = system.model.serialize()
to_save.update(train_set.get_infos())
torch.save(to_save, os.path.join(exp_dir, "best_model.pth"))
