def main(configfile, folddir, resultsdir, outputfile):
with open(configfile, 'r') as f:
configuration = yaml.load(f)
# Generate tasks from experiment list
tasks = {}
for experiment in configuration['experiments']:
classifier = experiment['classifier']
dataset = experiment['dataset']
folds = data.get_folds(folddir, dataset)
for f in range(len(folds)):
for r in range(experiment['reps']):
key = (classifier, dataset, experiment['kernel'], f, r)
task = Task(*key)
tasks[key] = task
# Mark finished tasks
for task in tasks.values():
predfile = os.path.join(resultsdir, task.filebase('preds'))
task.predfile = predfile
if os.path.exists(predfile):
task.finish()
reindexed = defaultdict(lambda: defaultdict(list))
for (c, d, k, f, r), task in tasks.items():
reindexed[(c, d, k)][r].append(task)
existing_keys = set()
if os.path.exists(outputfile):
with open(outputfile, 'r') as f:
for line in f:
c, d, k = line.strip().split(',')[:3]
existing_keys.add((c, d, k))
with open(outputfile, 'a+') as f:
rep_aucs = defaultdict(list)
for key, reps in sorted(reindexed.items()):
if key in existing_keys:
print 'Skipping %s (already finished)...' % str(key)
continue
data_dict = data.get_dataset(key[1])
bag_ids = sorted(data_dict.keys())
y_true = [data_dict[bid][1] for bid in bag_ids]
predictions = []
for rep, task_list in sorted(reps.items()):
if all(task.finished for task in task_list):
predictions.append(get_preds(key, task_list, bag_ids))
else:
break
if len(predictions) != len(reps):
print 'Skipping %s (incomplete)...' % str(key)
continue
predictions = np.vstack(predictions)
# We want a cumulative average, but doesn't matter for AUC
cumpreds = np.cumsum(predictions, axis=0)
aucs = [auc_score(y_true, cp) for cp in cumpreds]
line = ','.join(map(str, key) + map(str, aucs))
print line
f.write(line + '\n')
def main(configfile, folddir, resultsdir):
with open(configfile, 'r') as f:
configuration = yaml.load(f)
# Generate tasks from experiment list
tasks = {}
for experiment in configuration['experiments']:
classifier = experiment['classifier']
dataset = experiment['dataset']
folds = get_folds(folddir, dataset)
for f in range(len(folds)):
for r in range(experiment['reps'] + 1):
key = (classifier, dataset, experiment['kernel'], f, r)
kwargs = {}
kwargs['params'] = experiment['params']
task = Task(*key, **kwargs)
tasks[key] = task
# Mark finished tasks
for task in tasks.values():
predfile = os.path.join(resultsdir, task.filebase('preds'))
if os.path.exists(predfile):
task.finish()
def handle(key, task, submission):
if 'stats' in submission:
sfile = os.path.join(resultsdir, task.filebase('stats'))
with open(sfile, 'w+') as f:
f.write(
yaml.dump(submission['stats'], default_flow_style=False))
pfile = os.path.join(resultsdir, task.filebase('preds'))
with open(pfile, 'w+') as f:
f.write(yaml.dump(submission['preds'], default_flow_style=False))
server = ExperimentServer(tasks, render, handle)
cherrypy.config.update({
'server.socket_port': PORT,
'server.socket_host': '0.0.0.0'
})
cherrypy.quickstart(server)
def main(dataset, folddir, outputdir, reps=0):
data_dict = data.get_dataset(dataset)
folds = data.get_folds(folddir, dataset)
all_bag_ids = set(data_dict.keys())
progress = ProgressMonitor(total=reps * len(folds),
msg='Generating Replicates')
for f in range(len(folds)):
test = data.get_fold(folddir, dataset, f)
bag_ids = np.array(list(all_bag_ids - set(test)))
n = len(bag_ids)
for r in range(1, reps + 1):
rep_path = os.path.join(outputdir,
'%s_%04d_%06d.rep' % (dataset, f, r))
if not os.path.exists(rep_path):
sample = np.random.randint(n, size=n)
sampled_bags = bag_ids[sample]
with open(rep_path, 'w+') as ofile:
ofile.write('\n'.join([bid for bid in sampled_bags.flat]))
progress.increment()
def main(configfile, folddir, resultsdir):
with open(configfile, 'r') as f:
configuration = yaml.load(f)
# Generate tasks from experiment list
total = 0
actual = 0
prog = ProgressMonitor(total=len(configuration['experiments']),
msg='Computing noise')
for experiment in configuration['experiments']:
technique = experiment['technique']
classifier = experiment['classifier']
dataset = experiment['dataset']
ids, _, y = data.get_dataset(dataset)
y_dict = {}
for (bid, iid), yi in zip(ids, y):
y_dict[bid, iid] = yi
folds = data.get_folds(folddir, dataset)
for f in range(len(folds)):
for r in range(experiment['reps']):
for i in experiment['initial']:
for s in experiment['shuffled']:
labeled = setup_rep(technique, experiment['noise'],
dataset, f, r, i, s, folddir,
resultsdir)
pos_shuffled = get_positive_shuffled(labeled, i, s)
total += len(pos_shuffled)
actual += count_actual_positive(pos_shuffled, y_dict)
prog.increment()
if total > 0:
print 1 - (float(actual) / total)
if total > 0:
print 1 - (float(actual) / total)
def main(configfile, folddir, resultsdir):
with open(configfile, 'r') as f:
configuration = yaml.load(f)
# Count total experiments for progress monitor
exps = 0
for experiment in configuration['experiments']:
dataset = experiment['dataset']
folds = get_folds(folddir, dataset)
for f in range(len(folds)):
for r in range(experiment['reps']):
for n in experiment['noise']:
for s in experiment['shuffled']:
exps += 1
prog = ProgressMonitor(total=exps, msg='Generating Shuffled Bags')
# Generate tasks from experiment list
tasks = {}
for experiment in configuration['experiments']:
technique = experiment['technique']
classifier = experiment['classifier']
dataset = experiment['dataset']
folds = get_folds(folddir, dataset)
for f in range(len(folds)):
for r in range(experiment['reps']):
for n in experiment['noise']:
for s in experiment['shuffled']:
key = (technique, classifier,
dataset,
experiment['kernel'],
f, r, n, s)
kwargs = {}
kwargs['params'] = experiment['params']
kwargs['shuffled_bags'] = setup_rep(technique, dataset,
f, r, n, s,
folddir, resultsdir)
task = Task(*key, **kwargs)
tasks[key] = task
prog.increment()
# Mark finished tasks
for task in tasks.values():
predfile = os.path.join(resultsdir, task.filebase('preds'))
if os.path.exists(predfile):
task.finish()
def handle(key, task, submission):
if 'stats' in submission:
sfile = os.path.join(resultsdir, task.filebase('stats'))
with open(sfile, 'w+') as f:
f.write(yaml.dump(submission['stats'], default_flow_style=False))
pfile = os.path.join(resultsdir, task.filebase('preds'))
with open(pfile, 'w+') as f:
f.write(yaml.dump(submission['preds'], default_flow_style=False))
server = ExperimentServer(tasks, render, handle)
cherrypy.config.update({'server.socket_port': PORT,
'server.socket_host': '0.0.0.0'})
cherrypy.quickstart(server)
def main(configfile, folddir, resultsdir, outputfile):
with open(configfile, 'r') as f:
configuration = yaml.load(f)
# Generate tasks from experiment list
tasks = {}
for experiment in configuration['experiments']:
technique = experiment['technique']
classifier = experiment['classifier']
dataset = experiment['dataset']
folds = data.get_folds(folddir, dataset)
for f in range(len(folds)):
for r in range(experiment['reps']):
for i in experiment['initial']:
for s in experiment['shuffled']:
key = (technique, classifier, dataset,
experiment['kernel'], f, r, i, s,
experiment['queries'])
task = Task(*key)
tasks[key] = task
# Mark finished tasks
for task in tasks.values():
predfile = os.path.join(resultsdir, task.filebase('preds'))
task.predfile = predfile
if os.path.exists(predfile):
task.finish()
reindexed = defaultdict(lambda: defaultdict(list))
for (t, c, d, k, f, r, i, s, q), task in tasks.items():
reindexed[(t, c, d, k, i, s, q)][r].append(task)
existing_keys = set()
if os.path.exists(outputfile):
with open(outputfile, 'r') as f:
for line in f:
t, c, d, k, i, s, q = line.strip().split(',')[:7]
existing_keys.add((t, c, d, k, int(i), int(s), int(q)))
with open(outputfile, 'a+') as f:
rep_aucs = defaultdict(list)
for key, reps in sorted(reindexed.items()):
if key in existing_keys:
print 'Skipping %s (already finished)...' % str(key)
continue
ids, _, y = data.get_dataset(key[2])
y_dict = defaultdict(bool)
for (bid, iid), yi in zip(ids, y):
y_dict[bid] |= bool(yi)
aucs = []
for rep, task_list in sorted(reps.items()):
if all(task.finished for task in task_list):
aucs.append(calc_auc_score(key, task_list, y_dict))
else:
break
if len(aucs) != len(reps):
print 'Skipping %s (incomplete)...' % str(key)
continue
aucs = np.vstack(aucs)
avg_aucs = np.average(aucs, axis=0)
line = ','.join(map(str, key) + map(str, avg_aucs.flat))
print line
f.write(line + '\n')
def main(args):
os.environ['KMP_WARNINGS'] = '0'
torch.cuda.manual_seed_all(1)
np.random.seed(0)
print(args.model_name)
print(args.alpha)
# filter array
num_features = [
args.features * i
for i in range(1, args.levels + 2 + args.levels_without_sample)
]
# 確定 輸出大小
target_outputs = int(args.output_size * args.sr)
# 訓練才保存模型設定參數
# 設定teacher and student and student_for_backward 超參數
student_KD = Waveunet(args.channels,
num_features,
args.channels,
levels=args.levels,
encoder_kernel_size=args.encoder_kernel_size,
decoder_kernel_size=args.decoder_kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res)
KD_optimizer = Adam(params=student_KD.parameters(), lr=args.lr)
print(25 * '=' + 'model setting' + 25 * '=')
print('student_KD: ', student_KD.shapes)
if args.cuda:
student_KD = utils.DataParallel(student_KD)
print("move student_KD to gpu\n")
student_KD.cuda()
state = {"step": 0, "worse_epochs": 0, "epochs": 0, "best_pesq": -np.Inf}
if args.load_model is not None:
print("Continuing full model from checkpoint " + str(args.load_model))
state = utils.load_model(student_KD, KD_optimizer, args.load_model,
args.cuda)
dataset = get_folds(args.dataset_dir, args.outside_test)
log_dir, checkpoint_dir, result_dir = utils.mkdir_and_get_path(args)
# print(model)
if args.test is False:
writer = SummaryWriter(log_dir)
# set hypeparameter
# printing hypeparameters info
print(25 * '=' + 'printing hypeparameters info' + 25 * '=')
with open(os.path.join(log_dir, 'config.json'), 'w') as f:
json.dump(args.__dict__, f, indent=5)
print('saving commandline_args')
student_size = sum(p.numel() for p in student_KD.parameters())
print('student_parameter count: ', str(student_size))
if args.teacher_model is not None:
teacher_num_features = [
24 * i
for i in range(1, args.levels + 2 + args.levels_without_sample)
]
teacher_model = Waveunet(
args.channels,
teacher_num_features,
args.channels,
levels=args.levels,
encoder_kernel_size=args.encoder_kernel_size,
decoder_kernel_size=args.decoder_kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res)
if args.cuda:
teacher_model = utils.DataParallel(teacher_model)
teacher_model.cuda()
# print("move teacher to gpu\n")
student_size = sum(p.numel() for p in student_KD.parameters())
teacher_size = sum(p.numel() for p in teacher_model.parameters())
print('student_parameter count: ', str(student_size))
print('teacher_model_parameter count: ', str(teacher_size))
print(f'compression raito :{100*(student_size/teacher_size)}%')
if args.teacher_model is not None:
print("load teacher model" + str(args.teacher_model))
_ = utils.load_model(teacher_model, None, args.teacher_model,
args.cuda)
teacher_model.eval()
# If not data augmentation, at least crop targets to fit model output shape
crop_func = partial(crop, shapes=student_KD.shapes)
### DATASET
train_data = SeparationDataset(dataset,
"train",
args.sr,
args.channels,
student_KD.shapes,
False,
args.hdf_dir,
audio_transform=crop_func)
val_data = SeparationDataset(dataset,
"test",
args.sr,
args.channels,
student_KD.shapes,
False,
args.hdf_dir,
audio_transform=crop_func)
dataloader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
worker_init_fn=utils.worker_init_fn,
pin_memory=True)
# Set up the loss function
if args.loss == "L1":
criterion = nn.L1Loss()
elif args.loss == "L2":
criterion = nn.MSELoss()
else:
raise NotImplementedError("Couldn't find this loss!")
My_criterion = customLoss()
### TRAINING START
print('TRAINING START')
batch_num = (len(train_data) // args.batch_size)
while state["epochs"] < 100:
# if state["epochs"]<10:
# args.alpha=1
# else:
# args.alpha=0
# print('fix alpha:',args.alpha)
memory_alpha = []
print("epoch:" + str(state["epochs"]))
student_KD.train()
# monitor_value
avg_origin_loss = 0
with tqdm(total=len(dataloader)) as pbar:
for example_num, (x, targets) in enumerate(dataloader):
if args.cuda:
x = x.cuda()
targets = targets.cuda()
if args.teacher_model is not None:
# Set LR for this iteration
#print('base_model from KD')
utils.set_cyclic_lr(KD_optimizer, example_num,
len(train_data) // args.batch_size,
args.cycles, args.min_lr, args.lr)
_, avg_student_KD_loss = utils.compute_loss(
student_KD,
x,
targets,
criterion,
compute_grad=False)
KD_optimizer.zero_grad()
KD_outputs, KD_hard_loss, KD_loss, KD_soft_loss = utils.KD_compute_loss(
student_KD,
teacher_model,
x,
targets,
My_criterion,
alpha=args.alpha,
compute_grad=True,
KD_method=args.KD_method)
KD_optimizer.step()
# calculate backwarded model MSE
avg_origin_loss += avg_student_KD_loss / batch_num
# add to tensorboard
writer.add_scalar("KD_loss", KD_loss, state["step"])
writer.add_scalar("KD_hard_loss", KD_hard_loss,
state["step"])
writer.add_scalar("KD_soft_loss", KD_soft_loss,
state["step"])
else: # no KD training
utils.set_cyclic_lr(KD_optimizer, example_num,
len(train_data) // args.batch_size,
args.cycles, args.min_lr, args.lr)
KD_optimizer.zero_grad()
KD_outputs, KD_hard_loss = utils.compute_loss(
student_KD,
x,
targets,
nn.MSELoss(),
compute_grad=True)
KD_optimizer.step()
avg_origin_loss += KD_hard_loss / batch_num
writer.add_scalar("student_KD_loss", KD_hard_loss,
state["step"])
### save wav ####
if example_num % args.example_freq == 0:
input_centre = torch.mean(
x[0, :, student_KD.shapes["output_start_frame"]:
student_KD.shapes["output_end_frame"]],
0) # Stereo not supported for logs yet
writer.add_audio("input:",
input_centre,
state["step"],
sample_rate=args.sr)
writer.add_audio("pred:",
torch.mean(KD_outputs[0], 0),
state["step"],
sample_rate=args.sr)
writer.add_audio("target",
torch.mean(targets[0], 0),
state["step"],
sample_rate=args.sr)
state["step"] += 1
pbar.update(1)
# VALIDATE
val_loss, val_metrics = validate(args, student_KD, criterion,
val_data)
print("ori VALIDATION FINISHED: LOSS: " + str(val_loss))
writer.add_scalar("avg_origin_loss", avg_origin_loss,
state["epochs"])
writer.add_scalar("val_enhance_pesq", val_metrics[0],
state["epochs"])
writer.add_scalar("val_improve_pesq", val_metrics[1],
state["epochs"])
writer.add_scalar("val_enhance_stoi", val_metrics[2],
state["epochs"])
writer.add_scalar("val_improve_stoi", val_metrics[3],
state["epochs"])
writer.add_scalar("val_enhance_SISDR", val_metrics[4],
state["epochs"])
writer.add_scalar("val_improve_SISDR", val_metrics[5],
state["epochs"])
# writer.add_scalar("val_COPY_pesq",val_metrics_copy[0], state["epochs"])
writer.add_scalar("val_loss", val_loss, state["epochs"])
# Set up training state dict that will also be saved into checkpoints
checkpoint_path = os.path.join(
checkpoint_dir, "checkpoint_" + str(state["epochs"]))
if val_metrics[0] < state["best_pesq"]:
state["worse_epochs"] += 1
else:
print("MODEL IMPROVED ON VALIDATION SET!")
state["worse_epochs"] = 0
state["best_pesq"] = val_metrics[0]
state["best_checkpoint"] = checkpoint_path
# CHECKPOINT
print("Saving model...")
utils.save_model(student_KD, KD_optimizer, state, checkpoint_path)
print('dump alpha_memory')
with open(os.path.join(log_dir, 'alpha_' + str(state["epochs"])),
"wb") as fp: #Pickling
pickle.dump(memory_alpha, fp)
state["epochs"] += 1
writer.close()
info = args.model_name
path = os.path.join(result_dir, info)
else:
PATH = args.load_model.split("/")
info = PATH[-3] + "_" + PATH[-1]
if (args.outside_test == True):
info += "_outside_test"
print(info)
path = os.path.join(result_dir, info)
#### TESTING ####
# Test loss
print("TESTING")
# eval metrics
_ = utils.load_model(student_KD, KD_optimizer, state["best_checkpoint"],
args.cuda)
test_metrics = evaluate(args, dataset["test"], student_KD)
test_pesq = test_metrics['pesq']
test_stoi = test_metrics['stoi']
test_SISDR = test_metrics['SISDR']
test_noise = test_metrics['noise']
if not os.path.exists(path):
os.makedirs(path)
utils.save_result(test_pesq, path, "pesq")
utils.save_result(test_stoi, path, "stoi")
utils.save_result(test_SISDR, path, "SISDR")
utils.save_result(test_noise, path, "noise")
def main(args):
os.environ['KMP_WARNINGS'] = '0'
torch.cuda.manual_seed_all(1)
np.random.seed(0)
# filter array
num_features = [
args.features * i
for i in range(1, args.levels + 2 + args.levels_without_sample)
]
# 確定 輸出大小
target_outputs = int(args.output_size * args.sr)
# 訓練才保存模型設定參數
# 設定teacher and student and student_for_backward 超參數
student_KD = Waveunet(args.channels,
num_features,
args.channels,
levels=args.levels,
encoder_kernel_size=args.encoder_kernel_size,
decoder_kernel_size=args.decoder_kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res)
KD_optimizer = Adam(params=student_KD.parameters(), lr=args.lr)
print(25 * '=' + 'model setting' + 25 * '=')
print('student_KD: ', student_KD.shapes)
if args.cuda:
student_KD = utils.DataParallel(student_KD)
print("move student_KD to gpu\n")
student_KD.cuda()
state = {"step": 0, "worse_epochs": 0, "epochs": 0, "best_pesq": -np.Inf}
if args.load_model is not None:
print("Continuing full model from checkpoint " + str(args.load_model))
state = utils.load_model(student_KD, KD_optimizer, args.load_model,
args.cuda)
dataset = get_folds(args.dataset_dir, args.outside_test)
log_dir, checkpoint_dir, result_dir = utils.mkdir_and_get_path(args)
# print(model)
if args.test is False:
writer = SummaryWriter(log_dir)
# set hypeparameter
# printing hypeparameters info
with open(os.path.join(log_dir, 'config.json'), 'w') as f:
json.dump(args.__dict__, f, indent=5)
print('saving commandline_args')
if args.teacher_model is not None:
print(25 * '=' + 'printing hypeparameters info' + 25 * '=')
print(f'KD_method = {args.KD_method}')
teacher_num_features = [
24 * i
for i in range(1, args.levels + 2 + args.levels_without_sample)
]
teacher_model = Waveunet(
args.channels,
teacher_num_features,
args.channels,
levels=args.levels,
encoder_kernel_size=args.encoder_kernel_size,
decoder_kernel_size=args.decoder_kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res)
student_copy = Waveunet(
args.channels,
num_features,
args.channels,
levels=args.levels,
encoder_kernel_size=args.encoder_kernel_size,
decoder_kernel_size=args.decoder_kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res)
copy_optimizer = Adam(params=student_copy.parameters(), lr=args.lr)
student_copy2 = Waveunet(
args.channels,
num_features,
args.channels,
levels=args.levels,
encoder_kernel_size=args.encoder_kernel_size,
decoder_kernel_size=args.decoder_kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res)
copy2_optimizer = Adam(params=student_copy2.parameters(),
lr=args.lr)
policy_network = RL(n_inputs=2,
kernel_size=6,
stride=1,
conv_type=args.conv_type,
pool_size=4)
PG_optimizer = Adam(params=policy_network.parameters(),
lr=args.RL_lr)
if args.cuda:
teacher_model = utils.DataParallel(teacher_model)
policy_network = utils.DataParallel(policy_network)
student_copy = utils.DataParallel(student_copy)
student_copy2 = utils.DataParallel(student_copy2)
# print("move teacher to gpu\n")
teacher_model.cuda()
# print("student_copy to gpu\n")
student_copy.cuda()
# print("student_copy2 to gpu\n")
student_copy2.cuda()
# print("move policy_network to gpu\n")
policy_network.cuda()
student_size = sum(p.numel() for p in student_KD.parameters())
teacher_size = sum(p.numel() for p in teacher_model.parameters())
print('student_parameter count: ', str(student_size))
print('teacher_model_parameter count: ', str(teacher_size))
print('RL_parameter count: ',
str(sum(p.numel() for p in policy_network.parameters())))
print(f'compression raito :{100*(student_size/teacher_size)}%')
if args.teacher_model is not None:
print("load teacher model" + str(args.teacher_model))
_ = utils.load_model(teacher_model, None, args.teacher_model,
args.cuda)
teacher_model.eval()
if args.load_RL_model is not None:
print("Continuing full RL_model from checkpoint " +
str(args.load_RL_model))
_ = utils.load_model(policy_network, PG_optimizer,
args.load_RL_model, args.cuda)
# If not data augmentation, at least crop targets to fit model output shape
crop_func = partial(crop, shapes=student_KD.shapes)
### DATASET
train_data = SeparationDataset(dataset,
"train",
args.sr,
args.channels,
student_KD.shapes,
False,
args.hdf_dir,
audio_transform=crop_func)
val_data = SeparationDataset(dataset,
"test",
args.sr,
args.channels,
student_KD.shapes,
False,
args.hdf_dir,
audio_transform=crop_func)
dataloader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
worker_init_fn=utils.worker_init_fn,
pin_memory=True)
# Set up the loss function
if args.loss == "L1":
criterion = nn.L1Loss()
elif args.loss == "L2":
criterion = nn.MSELoss()
else:
raise NotImplementedError("Couldn't find this loss!")
My_criterion = customLoss()
### TRAINING START
print('TRAINING START')
if state["epochs"] > 0:
state["epochs"] = state["epochs"] + 1
batch_num = (len(train_data) // args.batch_size)
if args.teacher_model is not None:
counting = 0
PG_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer=PG_optimizer, gamma=args.decayRate)
while counting < state["epochs"]:
PG_optimizer.zero_grad()
PG_optimizer.step()
counting += 1
PG_lr_scheduler.step()
# print(f'modify lr RL rate : {counting} , until : {state["epochs"]}')
while state["epochs"] < 100:
memory_alpha = []
print("epoch:" + str(state["epochs"]))
# monitor_value
total_avg_reward = 0
total_avg_scalar_reward = 0
avg_origin_loss = 0
all_avg_KD_rate = 0
same = 0
with tqdm(total=len(dataloader)) as pbar:
for example_num, (x, targets) in enumerate(dataloader):
# if example_num==20:
# break
student_KD.train()
if args.cuda:
x = x.cuda()
targets = targets.cuda()
if args.teacher_model is not None:
student_copy.train()
student_copy2.train()
# Set LR for this iteration
temp = {'state_dict': None, 'optim_dict': None}
temp['state_dict'] = copy.deepcopy(
student_KD.state_dict())
temp['optim_dict'] = copy.deepcopy(
KD_optimizer.state_dict())
#print('base_model from KD')
student_KD.load_state_dict(temp['state_dict'])
KD_optimizer.load_state_dict(temp['optim_dict'])
student_copy.load_state_dict(temp['state_dict'])
copy_optimizer.load_state_dict(temp['optim_dict'])
student_copy2.load_state_dict(temp['state_dict'])
copy2_optimizer.load_state_dict(temp['optim_dict'])
utils.set_cyclic_lr(KD_optimizer, example_num,
len(train_data) // args.batch_size,
args.cycles, args.min_lr, args.lr)
utils.set_cyclic_lr(copy_optimizer, example_num,
len(train_data) // args.batch_size,
args.cycles, args.min_lr, args.lr)
utils.set_cyclic_lr(copy2_optimizer, example_num,
len(train_data) // args.batch_size,
args.cycles, args.min_lr, args.lr)
# forward student and teacher get output
student_KD_output, avg_student_KD_loss = utils.compute_loss(
student_KD,
x,
targets,
criterion,
compute_grad=False)
teacher_output, _ = utils.compute_loss(
teacher_model,
x,
targets,
criterion,
compute_grad=False)
# PG_state
diff_from_target = targets.detach(
) - student_KD_output.detach()
diff_from_teacher = teacher_output.detach(
) - student_KD_output.detach()
PG_state = torch.cat(
(diff_from_target, diff_from_teacher), 1)
# forward RL get alpha
alpha = policy_network(PG_state)
nograd_alpha = alpha.detach()
avg_KD_rate = torch.mean(nograd_alpha).item()
all_avg_KD_rate += avg_KD_rate / batch_num
KD_optimizer.zero_grad()
KD_outputs, KD_hard_loss, KD_loss, KD_soft_loss = utils.KD_compute_loss(
student_KD,
teacher_model,
x,
targets,
My_criterion,
alpha=nograd_alpha,
compute_grad=True,
KD_method=args.KD_method)
KD_optimizer.step()
copy_optimizer.zero_grad()
_, _, _, _ = utils.KD_compute_loss(
student_copy,
teacher_model,
x,
targets,
My_criterion,
alpha=1,
compute_grad=True,
KD_method=args.KD_method)
copy_optimizer.step()
copy2_optimizer.zero_grad()
_, _, _, _ = utils.KD_compute_loss(
student_copy2,
teacher_model,
x,
targets,
My_criterion,
alpha=0,
compute_grad=True,
KD_method=args.KD_method)
copy2_optimizer.step()
# calculate backwarded model MSE
backward_KD_loss = utils.loss_for_sample(
student_KD, x, targets)
backward_copy_loss = utils.loss_for_sample(
student_copy, x, targets)
backward_copy2_loss = utils.loss_for_sample(
student_copy2, x, targets)
# calculate rewards
rewards, same_num, before_decay = utils.get_rewards(
backward_KD_loss.detach(),
backward_copy_loss.detach(),
backward_copy2_loss.detach(),
backward_KD_loss.detach(), len(train_data),
state["epochs"] + 1)
same += same_num
rewards = rewards.detach()
avg_origin_loss += avg_student_KD_loss / batch_num
# avg_reward
avg_reward = torch.mean(rewards)
avg_scalar_reward = torch.mean(torch.abs(rewards))
total_avg_reward += avg_reward.item() / batch_num
total_avg_scalar_reward += avg_scalar_reward.item(
) / batch_num
# append to memory_alpha
nograd_alpha = nograd_alpha.detach().cpu()
memory_alpha.append(nograd_alpha.numpy())
PG_optimizer.zero_grad()
_ = utils.RL_compute_loss(alpha, rewards, nn.MSELoss())
PG_optimizer.step()
# print info
# print(f'avg_KD_rate = {avg_KD_rate} ')
# print(f'student_KD_loss = {avg_student_KD_loss}')
# print(f'backward_student_copy_loss = {np.mean(backward_copy_loss.detach().cpu().numpy())}')
# print(f'backward_student_KD_loss = {np.mean(backward_KD_loss.detach().cpu().numpy())}')
# print(f'backward_student_copy2_loss = {np.mean(backward_copy2_loss.detach().cpu().numpy())}')
# print(f'avg_reward = {avg_reward}')
# print(f'total_avg_reward = {total_avg_reward}')
# print(f'same = {same}')
# add to tensorboard
writer.add_scalar("student_KD_loss",
avg_student_KD_loss, state["step"])
writer.add_scalar(
"backward_student_KD_loss",
np.mean(backward_KD_loss.detach().cpu().numpy()),
state["step"])
writer.add_scalar("KD_loss", KD_loss, state["step"])
writer.add_scalar("KD_hard_loss", KD_hard_loss,
state["step"])
writer.add_scalar("KD_soft_loss", KD_soft_loss,
state["step"])
writer.add_scalar("avg_KD_rate", avg_KD_rate,
state["step"])
writer.add_scalar("rewards", avg_reward, state["step"])
writer.add_scalar("scalar_rewards", avg_scalar_reward,
state["step"])
writer.add_scalar("before_decay", before_decay,
state["step"])
else: # no KD training
utils.set_cyclic_lr(KD_optimizer, example_num,
len(train_data) // args.batch_size,
args.cycles, args.min_lr, args.lr)
KD_optimizer.zero_grad()
KD_outputs, KD_hard_loss = utils.compute_loss(
student_KD,
x,
targets,
nn.MSELoss(),
compute_grad=True)
KD_optimizer.step()
avg_origin_loss += KD_hard_loss / batch_num
writer.add_scalar("student_KD_loss", KD_hard_loss,
state["step"])
### save wav ####
if example_num % args.example_freq == 0:
input_centre = torch.mean(
x[0, :, student_KD.shapes["output_start_frame"]:
student_KD.shapes["output_end_frame"]],
0) # Stereo not supported for logs yet
# target=torch.mean(targets[0], 0).cpu().numpy()
# pred=torch.mean(KD_outputs[0], 0).detach().cpu().numpy()
# inputs=input_centre.cpu().numpy()
writer.add_audio("input:",
input_centre,
state["step"],
sample_rate=args.sr)
writer.add_audio("pred:",
torch.mean(KD_outputs[0], 0),
state["step"],
sample_rate=args.sr)
writer.add_audio("target",
torch.mean(targets[0], 0),
state["step"],
sample_rate=args.sr)
state["step"] += 1
pbar.update(1)
# VALIDATE
val_loss, val_metrics = validate(args, student_KD, criterion,
val_data)
print("ori VALIDATION FINISHED: LOSS: " + str(val_loss))
choose_val = val_metrics
if args.teacher_model is not None:
for i in range(len(nograd_alpha)):
writer.add_scalar("KD_rate_" + str(i), nograd_alpha[i],
state["epochs"])
print(f'all_avg_KD_rate = {all_avg_KD_rate}')
writer.add_scalar("all_avg_KD_rate", all_avg_KD_rate,
state["epochs"])
# writer.add_scalar("val_loss_copy", val_loss_copy, state["epochs"])
writer.add_scalar("total_avg_reward", total_avg_reward,
state["epochs"])
writer.add_scalar("total_avg_scalar_reward",
total_avg_scalar_reward, state["epochs"])
RL_checkpoint_path = os.path.join(
checkpoint_dir, "RL_checkpoint_" + str(state["epochs"]))
utils.save_model(policy_network, PG_optimizer, state,
RL_checkpoint_path)
PG_lr_scheduler.step()
writer.add_scalar("same", same, state["epochs"])
writer.add_scalar("avg_origin_loss", avg_origin_loss,
state["epochs"])
writer.add_scalar("val_enhance_pesq", choose_val[0],
state["epochs"])
writer.add_scalar("val_improve_pesq", choose_val[1],
state["epochs"])
writer.add_scalar("val_enhance_stoi", choose_val[2],
state["epochs"])
writer.add_scalar("val_improve_stoi", choose_val[3],
state["epochs"])
writer.add_scalar("val_enhance_SISDR", choose_val[4],
state["epochs"])
writer.add_scalar("val_improve_SISDR", choose_val[5],
state["epochs"])
# writer.add_scalar("val_COPY_pesq",val_metrics_copy[0], state["epochs"])
writer.add_scalar("val_loss", val_loss, state["epochs"])
# Set up training state dict that will also be saved into checkpoints
checkpoint_path = os.path.join(
checkpoint_dir, "checkpoint_" + str(state["epochs"]))
if choose_val[0] < state["best_pesq"]:
state["worse_epochs"] += 1
else:
print("MODEL IMPROVED ON VALIDATION SET!")
state["worse_epochs"] = 0
state["best_pesq"] = choose_val[0]
state["best_checkpoint"] = checkpoint_path
# CHECKPOINT
print("Saving model...")
utils.save_model(student_KD, KD_optimizer, state, checkpoint_path)
print('dump alpha_memory')
with open(os.path.join(log_dir, 'alpha_' + str(state["epochs"])),
"wb") as fp: #Pickling
pickle.dump(memory_alpha, fp)
state["epochs"] += 1
writer.close()
info = args.model_name
path = os.path.join(result_dir, info)
else:
PATH = args.load_model.split("/")
info = PATH[-3] + "_" + PATH[-1]
if (args.outside_test == True):
info += "_outside_test"
print(info)
path = os.path.join(result_dir, info)
# test_data = SeparationDataset(dataset, "test", args.sr, args.channels, student_KD.shapes, False, args.hdf_dir, audio_transform=crop_func)
#### TESTING ####
# Test loss
print("TESTING")
# eval metrics
#ling_data=get_ling_data_list('/media/hd03/sutsaiwei_data/data/mydata/ling_data')
#validate(args, student_KD, criterion, test_data)
#test_metrics = ling_evaluate(args, ling_data['noisy'], student_KD)
#test_metrics = evaluate_without_noisy(args, dataset["test"], student_KD)
test_metrics = evaluate(args, dataset["test"], student_KD)
test_pesq = test_metrics['pesq']
test_stoi = test_metrics['stoi']
test_SISDR = test_metrics['SISDR']
test_noise = test_metrics['noise']
if not os.path.exists(path):
os.makedirs(path)
utils.save_result(test_pesq, path, "pesq")
utils.save_result(test_stoi, path, "stoi")
utils.save_result(test_SISDR, path, "SISDR")
utils.save_result(test_noise, path, "noise")