def main(conf):
config = ParamConfig(
conf["training"]["batch_size"],
conf["training"]["epochs"],
conf["training"]["num_workers"],
cuda=True,
use_half=False,
learning_rate=conf["optim"]["lr"],
)
val_dataset = AVSpeechDataset(Path("data/val.csv"), Path(EMBED_DIR),
conf["main_args"]["n_src"])
model = load_best_model(conf, conf["main_args"]["exp_dir"])
print(
f"AVFusion has {sum(np.prod(i.shape) for i in model.parameters()):,} parameters"
)
if torch.cuda.device_count() > 1:
print(f"Multiple GPUs available")
device_ids = (list(map(int, conf["main_args"]["gpus"].split(",")))
if conf["main_args"]["gpus"] != "-1" else None)
model = torch.nn.DataParallel(model, device_ids=device_ids)
validate(model, val_dataset, config)
def main(conf):
# Get best trained model
model = load_best_model(conf["train_conf"], conf["exp_dir"])
if conf["use_gpu"]:
model = model.cuda()
# Evaluate performances separately w/ and w/o reverb
for subdir in ["with_reverb", "no_reverb"]:
dict_list = get_wavs_dict_list(os.path.join(conf["test_dir"], subdir))
save_dir = os.path.join(conf["exp_dir"], subdir + "examples/")
os.makedirs(save_dir, exist_ok=True)
all_metrics_df = evaluate(dict_list,
model,
conf=conf,
save_dir=save_dir)
all_metrics_df.to_csv(
os.path.join(conf["exp_dir"], "all_metrics_{}.csv".format(subdir)))
# Print and save summary metrics
final_results = {}
for metric_name in COMPUTE_METRICS:
input_metric_name = "input_" + metric_name
ldf = all_metrics_df[metric_name] - all_metrics_df[
input_metric_name]
final_results[metric_name] = all_metrics_df[metric_name].mean()
final_results[metric_name + "_imp"] = ldf.mean()
print("Overall metrics {} :".format(subdir))
pprint(final_results)
filename = os.path.join(conf["exp_dir"],
"final_metrics_{}.json".format(subdir))
with open(filename, "w") as f:
json.dump(final_results, f, indent=0)
def __init__(self):
self._komi = 0
self.board, self.player = game_init()
model = load_best_model()
self.sejong_engine = SejongGoEngine(model, conf['MCTS_SIMULATIONS'],
self.board)
print("GTP engine ready")
def load_model(self):
best_model = load_best_model()
logger.info("Loaded best model %s", best_model.name)
latest_model = load_latest_model()
logger.info("Loaded latest %s", latest_model.name)
return latest_model, best_model
def main(conf):
# Get best trained model
model = load_best_model(conf['train_conf'], conf['exp_dir'])
if conf['use_gpu']:
model = model.cuda()
model_device = next(model.parameters()).device
# Get a list of wav files (or single wav file)
save_folder = os.path.join(conf['exp_dir'], 'denoise')
os.makedirs(save_folder, exist_ok=True)
if os.path.isfile(conf['denoise_path']):
all_wavs = [conf['denoise_path']]
else:
# If this is a bunch of files we need to denoise, call the subdir
# of denoise the same way as the basename of the denoise dir.
save_folder = os.path.join(save_folder,
os.path.basename(conf['denoise_path']))
all_wavs = glob.glob(conf['denoise_path'] + '*.wav')
for wav_path in tqdm(all_wavs):
mix, fs = sf.read(wav_path, dtype='float32')
with torch.no_grad():
net_inp = torch.tensor(mix)[None].to(model_device)
estimate = model.denoise(net_inp).squeeze().cpu().data.numpy()
# Save the estimate speech
wav_name = os.path.basename(wav_path)
sf.write(os.path.join(save_folder, wav_name), estimate, fs)
def main(conf):
# Get best trained model
model = load_best_model(conf['train_conf'], conf['exp_dir'])
if conf['use_gpu']:
model = model.cuda()
# Evaluate performances separately w/ and w/o reverb
for subdir in ['with_reverb', 'no_reverb']:
dict_list = get_wavs_dict_list(os.path.join(conf['test_dir'], subdir))
save_dir = os.path.join(conf['exp_dir'], subdir + 'examples/')
os.makedirs(save_dir, exist_ok=True)
all_metrics_df = evaluate(dict_list, model, conf=conf,
save_dir=save_dir)
all_metrics_df.to_csv(os.path.join(conf['exp_dir'],
'all_metrics_{}.csv'.format(subdir)))
# Print and save summary metrics
final_results = {}
for metric_name in COMPUTE_METRICS:
input_metric_name = 'input_' + metric_name
ldf = all_metrics_df[metric_name] - all_metrics_df[
input_metric_name]
final_results[metric_name] = all_metrics_df[metric_name].mean()
final_results[metric_name + '_imp'] = ldf.mean()
print('Overall metrics {} :'.format(subdir))
pprint(final_results)
filename = os.path.join(conf['exp_dir'],
'final_metrics_{}.json'.format(subdir))
with open(filename, 'w') as f:
json.dump(final_results, f, indent=0)
def train_by_multi_gpus(n_gpu=1, epochs=None):
import tensorflow as tf
logger.info("Training with %s GPUs", n_gpu)
if n_gpu > 1:
with tf.device('/cpu:0'):
model = load_latest_model()
else:
model = load_latest_model()
with tf.device('/cpu:0'):
best_model = load_best_model()
base_name, index = model.name.split('_')
new_name = "_".join([base_name, str(int(index) + 1)]) + ".h5"
all_data_file_names = get_file_names_data_dir(os.path.join(SELF_PLAY_DATA, best_model.name))
tf_callback = TensorBoard(log_dir=os.path.join(conf['LOG_DIR'], new_name),
histogram_freq=conf['HISTOGRAM_FREQ'], batch_size=BATCH_SIZE, write_graph=False,
write_grads=False)
nan_callback = TerminateOnNaN()
if n_gpu > 1:
pmodel = multi_gpu_model(model, gpus=n_gpu)
opt = SGD(lr=1e-2, momentum=0.9)
pmodel.compile(loss=loss, optimizer=opt, metrics=["accuracy"])
else:
pmodel = model
if epochs is None:
epochs = EPOCHS_PER_SAVE
for epoch in tqdm.tqdm(range(epochs), desc="Epochs"):
for worker in tqdm.tqdm(range(NUM_WORKERS), desc="Iteration"):
files = sample(all_data_file_names, BATCH_SIZE) # RANDOM because we use SGD (Stochastic Gradient Decent)
X = np.zeros((BATCH_SIZE, SIZE, SIZE, 17))
policy_y = np.zeros((BATCH_SIZE, 1))
value_y = np.zeros((BATCH_SIZE, SIZE * SIZE + 1))
for j, filename in enumerate(files):
with h5py.File(filename) as f:
board = f['board'][:]
policy = f['policy_target'][:]
value_target = f['value_target'][()]
X[j] = board
policy_y[j] = value_target
value_y[j] = policy
fake_epoch = epoch * NUM_WORKERS + worker # used as initial_epoch, epochs is to be understood as "final epoch". The model is trained until the epoch of index epochs is reached.
pmodel.fit(X, [value_y, policy_y],
initial_epoch=fake_epoch,
epochs=fake_epoch + 1,
validation_split=VALIDATION_SPLIT, # Needed for TensorBoard histograms and gradi
callbacks=[tf_callback, nan_callback],
verbose=0, batch_size=BATCH_SIZE)
model.name = new_name.split('.')[0]
model.save(os.path.join(conf['MODEL_DIR'], new_name))
logger.info("Finished training with multi GPUs. New model %s saved", new_name)
return model
def main():
print("Starting run (v{})".format(__version__))
init_directories()
model_name = "model_1"
model = create_initial_model(name=model_name)
while True:
model = load_latest_model()
best_model = load_best_model()
train(model, game_model_name=best_model.name)
evaluate(best_model, model)
K.clear_session()
def main():
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
K.set_session(tf.Session(config=config))
init_directories()
while True:
model = load_latest_model()
best_model = load_best_model()
evaluate(best_model, model)
train(model, game_model_name=best_model.name)
K.clear_session()
def consumer(pipe):
model = load_best_model()
engine = Engine(model, 0)
output_p, input_p = pipe
input_p.close()
while True:
try:
item = output_p.recv()
result = engine.parse_command(item)
except EOFError:
break
print("<<<" + result)
print("consumer done")
def setUp(self):
init_directories()
model_name = "model_1"
model = create_initial_model(name=model_name)
best_model = load_best_model()
if best_model.name == model.name:
train(model, game_model_name=best_model.name)
evaluate(best_model, model)
# We save wether or not it was a better model
full_filename = os.path.join(conf['MODEL_DIR'], conf['BEST_MODEL'])
model.save(full_filename)
else:
model = best_model
self.model = model
def main():
model = load_best_model()
engine = Engine(model)
with open('logs/gtp.log', 'w') as f:
for line in sys.stdin:
f.write("<<<" + line)
result = engine.parse_command(line)
if result.strip():
sys.stdout.write(result)
sys.stdout.flush()
f.write("'''" + str(engine.player) + '\n')
f.write(">>>" + result)
f.flush()
def main():
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
K.set_session(tf.Session(config=config))
init_directories()
model_name = "model_1"
model = create_initial_model(name=model_name)
while True:
best_model = load_best_model()
self_play(best_model,
n_games=10,
mcts_simulations=conf['MCTS_SIMULATIONS'])
K.clear_session()
def main():
model = load_best_model()
gtp_log = os.path.join(conf['ROOT_DIR'], conf['LOG_DIR'], conf['GTP_LOGFILE'])
engine = Engine(model, gtp_log)
with open(gtp_log, 'w') as f:
for line in sys.stdin:
f.write("<<<" + line)
result = engine.parse_command(line)
if result.strip():
sys.stdout.write(result)
sys.stdout.flush()
f.write("'''" + str(engine.player) + '\n')
f.write(">>>" + result)
f.flush()
def main():
global start_time, time_limit
start_time = datetime.now()
time_limit = 11.6
raw = Utils.json_input()
request = raw['requests']
model = load_best_model()
board, player = game_init()
if not (request['x'] == -2 and request['y'] == -2):
board, player = make_play(request['x'] - 1, request['y'] - 1, board)
if player == 1:
color = 'B'
else:
color = 'W'
engine = ModelEngine(model, conf['MCTS_SIMULATIONS'], board)
x, y, _, _, _, _, _ = engine.genmove(color)
response = {}
response['x'], response['y'] = x + 1, y + 1
if y == conf['SIZE']:
response['x'], response['y'] = -1, -1
Utils.json_output({'response': response})
print(">>>BOTZONE_REQUEST_KEEP_RUNNING<<<")
time_limit = 3.6
while True:
try:
start_time = datetime.now()
raw = Utils.json_input()
request = raw['requests']
if not (request['x'] == -1 and request['y'] == -1):
engine.play('B',
request['x'] - 1,
request['y'] - 1,
update_tree=True)
x, y, _, _, _, _, _ = engine.genmove(color)
response = {}
response['x'], response['y'] = x + 1, y + 1
if y == conf['SIZE']:
response['x'], response['y'] = -1, -1
Utils.json_output({'response': response})
print(">>>BOTZONE_REQUEST_KEEP_RUNNING<<<")
except json.JSONDecodeError:
break
def run(target, start=0, iteration=10, batch_size=256, epochs=500, out='./model', name=None):
x_train, y_train, x_test = get_data(target=target, name=name)
print(f'Kind of Data: {len(x_train)}')
x_pred_list, y_pred_list, m_pred_list, v_pred_list = [], [], [], []
for i in range(iteration):
model = set_model(target, x_train, out=out, name=name)
hist = train(model, x_train, y_train,
epochs=epochs,
batch_size=batch_size,
label=LABEL[target],
seq=start+i,
out=out)
best_model = load_best_model(target, label=LABEL[target], seq=start+i, out=out)
pred_data_test = best_model.predict(x_test)
if target == 0:
x_pred_list.append(pred_data_test[:,0])
y_pred_list.append(pred_data_test[:,1])
if target == 1 or target == 3:
m_pred_list.append(pred_data_test[:,2])
if target == 2 or target == 3:
v_pred_list.append(pred_data_test[:,3])
if target == 0:
x_pred = np.mean(x_pred_list, axis=0)
y_pred = np.mean(y_pred_list, axis=0)
if target == 1 or target == 3:
m_pred = np.mean(m_pred_list, axis=0)
if target == 2 or target == 3:
v_pred = np.mean(v_pred_list, axis=0)
# submit
submit = pd.read_csv('xy_mlp_ensemble10_2nd.csv')
if target == 0:
submit.iloc[:, 1] = x_pred
submit.iloc[:, 2] = y_pred
if target == 1 or target == 3:
submit.iloc[:, 3] = m_pred
if target == 2 or target == 3:
submit.iloc[:, 4] = v_pred
submit.to_csv(os.path.join(out, 'mv_ensemble10.csv'), index=None)
def main():
model = load_latest_model()
best_model = load_best_model()
evaluate(best_model, model)
K.clear_session()
event_handler = MyHandler()
observer = Observer()
observer.schedule(event_handler, path=os.path.join(conf['MODEL_DIR']), recursive=False)
observer.start()
try:
while True:
time.sleep(60)
except KeyboardInterrupt:
observer.stop()
observer.join()
def main(conf):
config = ParamConfig(
conf["training"]["batch_size"],
conf["training"]["epochs"],
conf["training"]["num_workers"],
cuda=True,
use_half=False,
learning_rate=conf["optim"]["lr"],
)
val_dataset = AVSpeechDataset(Path("data/val.csv"), Path(EMBED_DIR),
conf["main_args"]["n_src"])
model = load_best_model(conf, conf["main_args"]["exp_dir"])
print(
f"AVFusion has {sum(np.prod(i.shape) for i in model.parameters()):,} parameters"
)
validate(model, val_dataset, config)
def main(conf):
set_trace()
test_set = WSJ2mixDataset(conf['data']['tt_wav_len_list'],
conf['data']['wav_base_path'] + '/tt',
sample_rate=conf['data']['sample_rate'])
test_loader = DataLoader(test_set,
shuffle=True,
batch_size=1,
num_workers=conf['data']['num_workers'],
drop_last=False)
istft = fb.Decoder(fb.STFTFB(**conf['filterbank']))
exp_dir = conf['main_args']['exp_dir']
model_path = os.path.join(exp_dir, 'checkpoints/_ckpt_epoch_0.ckpt')
model = load_best_model(conf, model_path)
pit_loss = PITLossWrapper(pairwise_mse, mode='pairwise')
system = DcSystem(model, None, None, None, config=conf)
# Randomly choose the indexes of sentences to save.
exp_dir = conf['main_args']['exp_dir']
exp_save_dir = os.path.join(exp_dir, 'examples/')
n_save = conf['main_args']['n_save_ex']
if n_save == -1:
n_save = len(test_set)
save_idx = random.sample(range(len(test_set)), n_save)
series_list = []
torch.no_grad().__enter__()
for batch in test_loader:
batch = [ele.type(torch.float32) for ele in batch]
inputs, targets, masks = system.unpack_data(batch)
est_targets = system(inputs)
mix_stft = system.enc(inputs.unsqueeze(1))
min_loss, min_idx = pit_loss.best_perm_from_perm_avg_loss(\
pairwise_mse, est_targets[1], masks)
for sidx in min_idx:
src_stft = mix_stft * est_targets[1][sidx]
src_sig = istft(src_stft)
def setUp(self):
init_directories()
model_name = "model_1"
best_model = load_best_model()
self.model = best_model
self.board, player = game_init()
def clear_board(self):
self.board, self.player = game_init()
return ""
def parse_command(self, line):
tokens = line.strip().split(" ")
command = tokens[0]
args = tokens[1:]
method = getattr(self, command)
result = method(*args)
if not result.strip():
return "=\n\n"
return "= " + result + "\n\n"
model = load_best_model()
engine = Engine(model, 0)
class SimpleEcho(WebSocket):
def handleMessage(self):
print("self.data: ", self.data)
result = engine.parse_command(self.data)
print("result: ", result)
self.sendMessage(result)
def handleConnected(self):
pass
def handleClose(self):
pass
def on_created(self, event):
time.sleep(30)
model = load_latest_model()
best_model = load_best_model()
evaluate(best_model, model)
K.clear_session()
def main(conf):
model = load_best_model(conf["train_conf"], conf["exp_dir"])
# Handle device placement
if conf["use_gpu"]:
model.cuda()
model_device = next(model.parameters()).device
test_set = WhamRDataset(
conf["test_dir"],
conf["task"],
sample_rate=conf["sample_rate"],
nondefault_nsrc=model.n_src,
segment=None,
) # Uses all segment length
# Used to reorder sources only
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx")
# Randomly choose the indexes of sentences to save.
ex_save_dir = os.path.join(conf["exp_dir"], "examples/")
if conf["n_save_ex"] == -1:
conf["n_save_ex"] = len(test_set)
save_idx = random.sample(range(len(test_set)), conf["n_save_ex"])
series_list = []
torch.no_grad().__enter__()
for idx in tqdm(range(len(test_set))):
# Forward the network on the mixture.
mix, sources = tensors_to_device(test_set[idx], device=model_device)
est_sources = model(mix[None, None])
loss, reordered_sources = loss_func(est_sources, sources[None], return_est=True)
mix_np = mix[None].cpu().data.numpy()
sources_np = sources.cpu().data.numpy()
est_sources_np = reordered_sources.squeeze(0).cpu().data.numpy()
utt_metrics = get_metrics(
mix_np,
sources_np,
est_sources_np,
sample_rate=conf["sample_rate"],
metrics_list=compute_metrics,
)
utt_metrics["mix_path"] = test_set.mix[idx][0]
series_list.append(pd.Series(utt_metrics))
# Save some examples in a folder. Wav files and metrics as text.
if idx in save_idx:
local_save_dir = os.path.join(ex_save_dir, "ex_{}/".format(idx))
os.makedirs(local_save_dir, exist_ok=True)
sf.write(local_save_dir + "mixture.wav", mix_np[0], conf["sample_rate"])
# Loop over the sources and estimates
for src_idx, src in enumerate(sources_np):
sf.write(local_save_dir + "s{}.wav".format(src_idx + 1), src, conf["sample_rate"])
for src_idx, est_src in enumerate(est_sources_np):
sf.write(
local_save_dir + "s{}_estimate.wav".format(src_idx + 1),
est_src,
conf["sample_rate"],
)
# Write local metrics to the example folder.
with open(local_save_dir + "metrics.json", "w") as f:
json.dump(utt_metrics, f, indent=0)
# Save all metrics to the experiment folder.
all_metrics_df = pd.DataFrame(series_list)
all_metrics_df.to_csv(os.path.join(conf["exp_dir"], "all_metrics.csv"))
# Print and save summary metrics
final_results = {}
for metric_name in compute_metrics:
input_metric_name = "input_" + metric_name
ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name]
final_results[metric_name] = all_metrics_df[metric_name].mean()
final_results[metric_name + "_imp"] = ldf.mean()
print("Overall metrics :")
pprint(final_results)
with open(os.path.join(conf["exp_dir"], "final_metrics.json"), "w") as f:
json.dump(final_results, f, indent=0)
def load_model(self):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id)
self.best_model = load_best_model()
self.latest_model = load_latest_model()
import model model = model.load_best_model()
def main(conf):
model = load_best_model(conf['train_conf'], conf['exp_dir'])
# Handle device placement
if conf['use_gpu']:
model.cuda()
model_device = next(model.parameters()).device
test_set = Wsj0mixDataset(conf['test_dir'],
n_src=conf['n_src'],
segment=None)
# Used to reorder sources only
loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from='pw_mtx')
# Randomly choose the indexes of sentences to save.
ex_save_dir = os.path.join(conf['exp_dir'], 'examples/')
if conf['n_save_ex'] == -1:
conf['n_save_ex'] = len(test_set)
save_idx = random.sample(range(len(test_set)), conf['n_save_ex'])
series_list = []
torch.no_grad().__enter__()
for idx in tqdm(range(len(test_set))):
# Forward the network on the mixture.
mix, sources = tensors_to_device(test_set[idx], device=model_device)
if conf['train_conf']['training']['loss_alpha'] == 1:
# If Deep clustering only, use DC masks.
est_sources, dic_out = model.dc_head_separate(mix[None, None])
else:
# If Chimera, use mask-inference head masks
est_sources, dic_out = model.separate(mix[None, None])
loss, reordered_sources = loss_func(est_sources,
sources[None],
return_est=True)
mix_np = mix[None].cpu().data.numpy()
sources_np = sources.cpu().data.numpy()
est_sources_np = reordered_sources.squeeze(0).cpu().data.numpy()
utt_metrics = get_metrics(mix_np,
sources_np,
est_sources_np,
sample_rate=conf['sample_rate'],
metrics_list=compute_metrics)
utt_metrics['mix_path'] = test_set.mix[idx][0]
series_list.append(pd.Series(utt_metrics))
# Save some examples in a folder. Wav files and metrics as text.
if idx in save_idx:
local_save_dir = os.path.join(ex_save_dir, 'ex_{}/'.format(idx))
os.makedirs(local_save_dir, exist_ok=True)
sf.write(local_save_dir + "mixture.wav", mix_np[0],
conf['sample_rate'])
# Loop over the sources and estimates
for src_idx, src in enumerate(sources_np):
sf.write(local_save_dir + "s{}.wav".format(src_idx + 1), src,
conf['sample_rate'])
for src_idx, est_src in enumerate(est_sources_np):
sf.write(
local_save_dir + "s{}_estimate.wav".format(src_idx + 1),
est_src, conf['sample_rate'])
# Write local metrics to the example folder.
with open(local_save_dir + 'metrics.json', 'w') as f:
json.dump(utt_metrics, f, indent=0)
# Save all metrics to the experiment folder.
all_metrics_df = pd.DataFrame(series_list)
all_metrics_df.to_csv(os.path.join(conf['exp_dir'], 'all_metrics.csv'))
# Print and save summary metrics
final_results = {}
for metric_name in compute_metrics:
input_metric_name = 'input_' + metric_name
ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name]
final_results[metric_name] = all_metrics_df[metric_name].mean()
final_results[metric_name + '_imp'] = ldf.mean()
print('Overall metrics :')
pprint(final_results)
with open(os.path.join(conf['exp_dir'], 'final_metrics.json'), 'w') as f:
json.dump(final_results, f, indent=0)
def main(conf):
best_model_path = os.path.join(conf["exp_dir"], "best_model.pth")
if not os.path.exists(best_model_path):
# make pth from checkpoint
model = load_best_model(conf["train_conf"],
conf["exp_dir"],
sample_rate=conf["sample_rate"])
torch.save(model.state_dict(), best_model_path)
else:
model, _ = make_model_and_optimizer(conf["train_conf"],
sample_rate=conf["sample_rate"])
model.eval()
model.load_state_dict(torch.load(best_model_path))
# Handle device placement
if conf["use_gpu"]:
model.cuda()
model_device = next(model.parameters()).device
test_dirs = [
conf["test_dir"].format(n_src)
for n_src in conf["train_conf"]["masknet"]["n_srcs"]
]
test_set = Wsj0mixVariable(
json_dirs=test_dirs,
n_srcs=conf["train_conf"]["masknet"]["n_srcs"],
sample_rate=conf["train_conf"]["data"]["sample_rate"],
seglen=None,
minlen=None,
)
# Randomly choose the indexes of sentences to save.
ex_save_dir = os.path.join(conf["exp_dir"], "examples/")
if conf["n_save_ex"] == -1:
conf["n_save_ex"] = len(test_set)
save_idx = random.sample(range(len(test_set)), conf["n_save_ex"])
series_list = []
torch.no_grad().__enter__()
for idx in tqdm(range(len(test_set))):
# Forward the network on the mixture.
mix, sources = [
torch.Tensor(x)
for x in tensors_to_device(test_set[idx], device=model_device)
]
est_sources = model.separate(mix[None])
p_si_snr = Penalized_PIT_Wrapper(pairwise_neg_sisdr_loss)(est_sources,
sources)
utt_metrics = {
"P-Si-SNR": p_si_snr.item(),
"counting_accuracy": float(sources.size(0) == est_sources.size(0)),
}
utt_metrics["mix_path"] = test_set.data[idx][0]
series_list.append(pd.Series(utt_metrics))
# Save some examples in a folder. Wav files and metrics as text.
if idx in save_idx:
mix_np = mix[None].cpu().data.numpy()
sources_np = sources.cpu().data.numpy()
est_sources_np = est_sources.cpu().data.numpy()
local_save_dir = os.path.join(ex_save_dir, "ex_{}/".format(idx))
os.makedirs(local_save_dir, exist_ok=True)
sf.write(local_save_dir + "mixture.wav", mix_np[0],
conf["sample_rate"])
# Loop over the sources and estimates
for src_idx, src in enumerate(sources_np):
sf.write(local_save_dir + "s{}.wav".format(src_idx + 1), src,
conf["sample_rate"])
for src_idx, est_src in enumerate(est_sources_np):
sf.write(
local_save_dir + "s{}_estimate.wav".format(src_idx + 1),
est_src,
conf["sample_rate"],
)
# Write local metrics to the example folder.
with open(local_save_dir + "metrics.json", "w") as f:
json.dump(utt_metrics, f, indent=0)
# Save all metrics to the experiment folder.
all_metrics_df = pd.DataFrame(series_list)
all_metrics_df.to_csv(os.path.join(conf["exp_dir"], "all_metrics.csv"))
# Print and save summary metrics
final_results = {}
for metric_name in ["P-Si-SNR", "counting_accuracy"]:
final_results[metric_name] = all_metrics_df[metric_name].mean()
print("Overall metrics :")
pprint(final_results)
with open(os.path.join(conf["exp_dir"], "final_metrics.json"), "w") as f:
json.dump(final_results, f, indent=0)
def load_model(self):
self.best_model = load_best_model()
self.latest_model = load_latest_model()
