def forward(model, generate_func, cuda, return_target):
"""Forward data to a model.
Args:
generate_func: generate function
cuda: bool
return_target: bool
Returns:
dict, keys: 'audio_name', 'output'; optional keys: 'target'
"""
outputs = []
audio_names = []
if return_target:
targets = []
# Evaluate on mini-batch
for data in generate_func:
if return_target:
(batch_x_for_an_audio, y, audio_name) = data
else:
(batch_x_for_an_audio, audio_name) = data
batch_x_for_an_audio = move_data_to_gpu(batch_x_for_an_audio, cuda)
# Predict
model.eval()
outputs_for_an_audio = model(batch_x_for_an_audio)
# logging.info('{}'.format(outputs_for_an_audio.data.cpu().numpy().shape))
# logging.info('{}'.format(audio_name))
# Append data
outputs.append(outputs_for_an_audio.data.cpu().numpy())
audio_names.append(audio_name)
if return_target:
targets.append(y)
dict = {}
outputs = np.array(outputs)
dict['output'] = outputs
audio_names = np.array(audio_names)
dict['audio_name'] = audio_names
if return_target:
targets = np.array(targets)
dict['target'] = targets
return dict
def forward_heatmap(model, generate_func, cuda, has_target):
"""Forward data to a model.
Args:
generate_func: generate function
cuda: bool
has_target: bool
Returns:
(outputs, targets, audio_names) | (outputs, audio_names)
"""
model.eval()
outputs = []
targets = []
audio_names = []
outputs_heatmap = [] ###############################
# Evaluate on mini-batch
for data in generate_func:
if has_target:
(batch_x, batch_y, batch_audio_names) = data
targets.append(batch_y)
else:
(batch_x, batch_audio_names) = data
batch_x = move_data_to_gpu(batch_x, cuda)
# Predict
batch_output, batch_heatmap = model(
batch_x) ###########################################
outputs.append(batch_output.data.cpu().numpy())
audio_names.append(batch_audio_names)
outputs_heatmap.append(batch_heatmap.data.cpu().numpy()
) #######################################
outputs = np.concatenate(outputs, axis=0)
audio_names = np.concatenate(audio_names, axis=0)
outputs_heatmap = np.concatenate(
outputs_heatmap, axis=0) ###########################################
if has_target:
targets = np.concatenate(targets, axis=0)
return outputs, targets, audio_names, outputs_heatmap ############################################
else:
return outputs, audio_names
def train(args):
# Arugments & parameters
dataset_dir = args.dataset_dir
workspace = args.workspace
validate = args.validate
holdout_fold = args.holdout_fold
mini_data = args.mini_data
cuda = args.cuda
filename = args.filename
classes_num = len(config.labels)
max_validate_iteration = 100
# Use partial data to speed up training
if mini_data:
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'mini_development.h5')
else:
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'development.h5')
if validate:
train_txt = os.path.join(dataset_dir, 'evaluation_setup',
'fold{}_train.txt'.format(holdout_fold))
validate_txt = os.path.join(dataset_dir, 'evaluation_setup',
'fold{}_evaluate.txt'.format(holdout_fold))
else:
train_txt = None
validate_txt = None
if validate:
models_dir = os.path.join(workspace, 'models', filename,
'holdout_fold{}'.format(holdout_fold))
else:
models_dir = os.path.join(workspace, 'models', filename, 'full_train')
create_folder(models_dir)
# Model
model = Model(classes_num)
if cuda:
model.cuda()
generator = DataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
train_txt=train_txt,
validate_txt=validate_txt)
# Optimizer
optimizer = optim.Adam(model.parameters(),
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.)
train_bgn_time = time.time()
# Train on mini-batch
for (iteration, (batch_x,
batch_y)) in enumerate(generator.generate_train()):
# Evaluate
if iteration % 200 == 0:
train_fin_time = time.time()
(tr_acc, tr_f1_score,
tr_loss) = evaluate(model=model,
generator=generator,
data_type='train',
max_iteration=max_validate_iteration,
cuda=cuda)
if validate:
(va_acc, va_f1_score,
va_loss) = evaluate(model=model,
generator=generator,
data_type='validate',
max_iteration=max_validate_iteration,
cuda=cuda)
train_time = train_fin_time - train_bgn_time
validate_time = time.time() - train_fin_time
# Print info
logging.info(
'iteration: {}, train time: {:.4f} s, validate time: {:.4f} s'
''.format(iteration, train_time, validate_time))
logging.info('tr_acc: {:.4f}, tr_f1_score: {:.4f}, tr_loss: {:.4f}'
''.format(tr_acc, tr_f1_score, tr_loss))
if validate:
logging.info('va_acc: {:.4f}, va_f1_score: {:.4f}, '
'va_loss: {:.4f}'.format(va_acc, va_f1_score,
va_loss))
logging.info('------------------------------------')
train_bgn_time = time.time()
# Save model
if iteration % 1000 == 0 and iteration > 0:
save_out_dict = {
'iteration': iteration,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_out_path = os.path.join(models_dir,
'md_{}_iters.tar'.format(iteration))
torch.save(save_out_dict, save_out_path)
logging.info('Model saved to {}'.format(save_out_path))
# Reduce learning rate
if iteration % 100 == 0 and iteration > 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.9
# Move data to gpu
batch_x = move_data_to_gpu(batch_x, cuda)
batch_y = move_data_to_gpu(batch_y, cuda)
# Train
model.train()
output = model(batch_x)
loss = F.nll_loss(output, batch_y)
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Stop learning
if iteration == 10000:
break
def forward(model, model_adv, generate_func, cuda, return_target):
"""Forward data to a model.
Args:
generate_func: generate function
cuda: bool
return_target: bool
Returns:
dict, keys: 'audio_name', 'output'; optional keys: 'target'
"""
outputs = []
audio_names = []
outputs_adv = []
if return_target:
targets = []
# Evaluate on mini-batch
for data in generate_func:
if return_target:
(batch_x, batch_y, batch_audio_names) = data
else:
(batch_x, batch_audio_names) = data
batch_x = move_data_to_gpu(batch_x, cuda)
# Predict
model.eval()
batch_output, _ = model(batch_x)
# advesarial predict
batch_y_pred = np.argmax(batch_output.data.cpu().numpy(), axis=-1)
batch_y_pred = move_data_to_gpu(batch_y_pred, cuda)
model_cp = copy.deepcopy(model)
for p in model_cp.parameters():
p.requires_grad = False
model_cp.eval()
model_adv.model = model_cp
del model_cp
batch_x_adv = model_adv.perturb(batch_x.data.cpu().numpy(),
batch_y_pred.data.cpu().numpy(),
cuda=cuda)
batch_x_adv = move_data_to_gpu(batch_x_adv, cuda)
batch_output_adv, _ = model(batch_x_adv)
# Append data
outputs.append(batch_output.data.cpu().numpy())
audio_names.append(batch_audio_names)
outputs_adv.append(batch_output_adv.data.cpu().numpy())
if return_target:
targets.append(batch_y)
dict = {}
outputs = np.concatenate(outputs, axis=0)
dict['output'] = outputs
audio_names = np.concatenate(audio_names, axis=0)
dict['audio_name'] = audio_names
outputs_adv = np.concatenate(outputs_adv, axis=0)
dict['output_adv'] = outputs_adv
if return_target:
targets = np.concatenate(targets, axis=0)
dict['target'] = targets
return dict
def train(args):
# Arugments & parameters
dataset_dir = args.dataset_dir
subdir = args.subdir
workspace = args.workspace
feature_type = args.feature_type
filename = args.filename
validation = args.validation
holdout_fold = args.holdout_fold
epsilon_value = args.epsilon_value
alpha_value = args.alpha_value
mini_data = args.mini_data
cuda = args.cuda
labels = config.labels
if 'mobile' in subdir:
devices = ['a', 'b', 'c']
else:
devices = ['a']
classes_num = len(labels)
# Paths
if mini_data:
hdf5_path = os.path.join(workspace, 'features', feature_type,
'mini_development.h5')
else:
hdf5_path = os.path.join(workspace, 'features', feature_type,
'development.h5')
if validation:
dev_train_csv = os.path.join(dataset_dir, subdir, 'evaluation_setup',
'fold{}_train.txt'.format(holdout_fold))
dev_validate_csv = os.path.join(
dataset_dir, subdir, 'evaluation_setup',
'fold{}_devel.txt'.format(holdout_fold))
models_dir = os.path.join(
workspace, 'models', subdir, filename,
'holdout_fold={}'.format(holdout_fold),
'epsilon={}-alpha={}'.format(epsilon_value, alpha_value))
else:
dev_train_csv = os.path.join(
dataset_dir, subdir, 'evaluation_setup',
'fold{}_traindevel.txt'.format(holdout_fold))
dev_validate_csv = os.path.join(dataset_dir, subdir,
'evaluation_setup',
'fold{}_test.txt'.format(holdout_fold))
models_dir = os.path.join(
workspace, 'models', subdir, filename, 'full_train',
'epsilon={}-alpha={}'.format(epsilon_value, alpha_value))
create_folder(models_dir)
# Model
model = Model(classes_num)
adversary = FGSMAttack(epsilon=epsilon_value, alpha=alpha_value)
if cuda:
model.cuda()
# Data generator
generator = DataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
dev_train_csv=dev_train_csv,
dev_validate_csv=dev_validate_csv)
# Optimizer
lr = 1e-3
optimizer = optim.Adam(model.parameters(),
lr=lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.)
train_bgn_time = time.time()
# Train on mini batches
for (iteration, (batch_x,
batch_y)) in enumerate(generator.generate_train()):
# Evaluate
if iteration % 100 == 0:
train_fin_time = time.time()
(tr_acc, tr_loss, tr_acc_adv,
tr_loss_adv) = evaluate(model=model,
model_adv=adversary,
generator=generator,
data_type='train',
devices=devices,
max_iteration=None,
cuda=cuda)
logging.info(
'tr_acc: {:.3f}, tr_loss: {:.3f}, tr_acc_adv: {:.3f}, tr_loss_adv: {:.3f}'
.format(tr_acc, tr_loss, tr_acc_adv, tr_loss_adv))
(va_acc, va_loss, va_acc_adv,
va_loss_adv) = evaluate(model=model,
model_adv=adversary,
generator=generator,
data_type='validate',
devices=devices,
max_iteration=None,
cuda=cuda)
logging.info(
'va_acc: {:.3f}, va_loss: {:.3f}, va_acc_adv: {:.3f}, va_loss_adv: {:.3f}'
.format(va_acc, va_loss, va_acc_adv, va_loss_adv))
train_time = train_fin_time - train_bgn_time
validate_time = time.time() - train_fin_time
logging.info(
'iteration: {}, train time: {:.3f} s, validate time: {:.3f} s'
''.format(iteration, train_time, validate_time))
logging.info('------------------------------------')
train_bgn_time = time.time()
# Save model
if iteration % 1000 == 0 and iteration > 0:
save_out_dict = {
'iteration': iteration,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_out_path = os.path.join(models_dir,
'md_{}_iters.tar'.format(iteration))
torch.save(save_out_dict, save_out_path)
logging.info('Model saved to {}'.format(save_out_path))
# Reduce learning rate
if iteration % 100 == 0 > 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.9
# Train
batch_x = move_data_to_gpu(batch_x, cuda)
batch_y = move_data_to_gpu(batch_y, cuda)
model.train()
batch_output, batch_outputvector = model(batch_x)
loss = F.nll_loss(batch_output, batch_y)
if iteration >= 1000:
batch_y_pred = np.argmax(batch_output.data.cpu().numpy(), axis=-1)
batch_y_pred = move_data_to_gpu(batch_y_pred, cuda)
model_cp = copy.deepcopy(model)
for p in model_cp.parameters():
p.requires_grad = False
model_cp.eval()
adversary.model = model_cp
del model_cp
batch_x_adv = adversary.perturb(batch_x.data.cpu().numpy(),
batch_y_pred.data.cpu().numpy(),
cuda=cuda)
batch_x_adv = move_data_to_gpu(batch_x_adv, cuda)
batch_output_adv, batch_outputvector_adv = model(batch_x_adv)
loss_adv = F.nll_loss(batch_output_adv, batch_y)
loss_pair = F.mse_loss(batch_outputvector_adv, batch_outputvector)
#print('loss:'+str(loss)+'\t'+'loss_adv'+str(loss_adv)+'\t'+'loss_pair'+str(loss_pair))
loss = 0.4 * loss + 0.4 * loss_adv + 0.2 * loss_pair
#if iteration % 10 == 0:
# logging.info('batch loss: {}, batch loss_adv: {}'.format(loss, loss_adv))
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Stop learning
if iteration == 10000:
break
def inference(args):
# Arugments & parameters
workspace = args.workspace
model_type = args.model_type
holdout_fold = args.holdout_fold
scene_type = args.scene_type
snr = args.snr
iteration = args.iteration
filename = args.filename
cuda = args.cuda
labels = config.labels
classes_num = len(labels)
sample_rate = config.sample_rate
window_size = config.window_size
overlap = config.overlap
hop_size = window_size - overlap
mel_bins = config.mel_bins
seq_len = config.seq_len
ix_to_lb = config.ix_to_lb
threshold = 0.1
# Paths
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'scene_type={},snr={}'.format(scene_type, snr),
'development.h5')
model_path = os.path.join(
workspace, 'models', filename, 'model_type={}'.format(model_type),
'scene_type={},snr={}'
''.format(scene_type, snr), 'holdout_fold{}'.format(holdout_fold),
'md_{}_iters.tar'.format(iteration))
yaml_path = os.path.join(workspace, 'mixture.yaml')
out_stat_path = os.path.join(
workspace, 'stats', filename, 'model_type={}'.format(model_type),
'scene_type={},snr={}'
''.format(scene_type,
snr), 'holdout_fold{}'.format(holdout_fold), 'stat.p')
create_folder(os.path.dirname(out_stat_path))
pred_prob_path = os.path.join(
workspace, 'pred_probs', filename, 'model_type={}'.format(model_type),
'scene_type={},snr={}'
''.format(scene_type,
snr), 'holdout_fold{}'.format(holdout_fold), 'pred_prob.p')
create_folder(os.path.dirname(pred_prob_path))
# Load yaml file
load_yaml_time = time.time()
with open(yaml_path, 'r') as f:
meta = yaml.load(f)
logging.info('Load yaml file time: {:.3f} s'.format(time.time() -
load_yaml_time))
feature_extractor = LogMelExtractor(sample_rate=sample_rate,
window_size=window_size,
overlap=overlap,
mel_bins=mel_bins)
# Load model
Model = get_model(model_type)
model = Model(classes_num, seq_len, mel_bins, cuda)
checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint['state_dict'])
if cuda:
model.cuda()
# Data generator
generator = InferenceDataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
holdout_fold=holdout_fold)
generate_func = generator.generate_validate(data_type='validate',
shuffle=False,
max_iteration=None)
audio_names = []
at_outputs = []
at_targets = []
sed_outputs = []
sed_targets = []
ss_outputs = []
ss_targets = []
validate_num = len(generator.validate_audio_indexes)
# Evaluate on mini-batch
for iteration, data in enumerate(generate_func):
print('{} / {} inferenced & detected!'.format(iteration * batch_size,
validate_num))
(batch_x, batch_y, batch_audio_names) = data
batch_x = move_data_to_gpu(batch_x, cuda)
# Predict
with torch.no_grad():
model.eval()
(batch_output, batch_bottleneck) = model(batch_x,
return_bottleneck=True)
batch_output = batch_output.data.cpu().numpy()
'''(batch_size, classes_num)'''
batch_bottleneck = batch_bottleneck.data.cpu().numpy()
'''(batch_size, classes_num, seq_len, mel_bins)'''
audio_names.append(batch_audio_names)
at_outputs.append(batch_output)
at_targets.append(batch_y)
batch_pred_sed = np.mean(batch_bottleneck, axis=-1)
batch_pred_sed = np.transpose(batch_pred_sed, (0, 2, 1))
'''(batch_size, seq_len, classes_num)'''
for n in range(len(batch_audio_names)):
gt_meta = search_meta_by_mixture_name(meta, batch_audio_names[n])
gt_events = gt_meta['events']
gt_sed = get_sed_from_meta(gt_events)
'''(seq_len, classes_num)'''
# Do audio tagging first, then only apply SED to the positive
# classes to reduce the false positives.
pred_classes = np.where(batch_output[n] > threshold)[0]
pred_sed = np.zeros((seq_len, classes_num))
pred_sed[:, pred_classes] = batch_pred_sed[n][:, pred_classes]
'''(seq_len, classes_num)'''
sed_outputs.append(pred_sed)
sed_targets.append(gt_sed)
(events_stft, scene_stft, mixture_stft) = \
generator.get_events_scene_mixture_stft(batch_audio_names[n])
'''(seq_len, fft_bins)'''
events_stft = np.dot(events_stft, feature_extractor.melW)
scene_stft = np.dot(scene_stft, feature_extractor.melW)
'''(seq_len, mel_bins)'''
gt_mask = ideal_binary_mask(events_stft, scene_stft)
'''(seq_len, mel_bins)'''
gt_masks = gt_mask[:, :, None] * gt_sed[:, None, :]
gt_masks = gt_masks.astype(np.float32)
'''(seq_len, fft_size, classes_num)'''
pred_masks = batch_bottleneck[n].transpose(1, 2, 0)
'''(seq_len, fft_size, classes_num)'''
ss_outputs.append(pred_masks)
ss_targets.append(gt_masks)
# if iteration == 3: break
audio_names = np.concatenate(audio_names, axis=0)
at_outputs = np.concatenate(at_outputs, axis=0)
at_targets = np.concatenate(at_targets, axis=0)
'''(audio_clips,)'''
sed_outputs = np.array(sed_outputs)
sed_targets = np.array(sed_targets)
'''(audio_clips, seq_len, classes_num)'''
ss_outputs = np.array(ss_outputs)
ss_targets = np.array(ss_targets)
'''(audio_clips, seq_len, mel_bins, classes_num)'''
pred_prob = {
'audio_name': audio_names,
'at_output': at_outputs,
'at_target': at_targets,
'sed_output': sed_outputs,
'sed_target': sed_targets
}
pickle.dump(pred_prob, open(pred_prob_path, 'wb'))
logging.info('Saved stat to {}'.format(pred_prob_path))
# Evaluate audio tagging
at_time = time.time()
(at_precision, at_recall,
at_f1_score) = prec_recall_fvalue(at_targets, at_outputs, threshold, None)
at_auc = metrics.roc_auc_score(at_targets, at_outputs, average=None)
at_ap = metrics.average_precision_score(at_targets,
at_outputs,
average=None)
logging.info('Audio tagging time: {:.3f} s'.format(time.time() - at_time))
# Evaluate SED
sed_time = time.time()
(sed_precision, sed_recall, sed_f1_score) = prec_recall_fvalue(
sed_targets.reshape((sed_targets.shape[0] * sed_targets.shape[1],
sed_targets.shape[2])),
sed_outputs.reshape((sed_outputs.shape[0] * sed_outputs.shape[1],
sed_outputs.shape[2])),
thres=threshold,
average=None)
sed_auc = metrics.roc_auc_score(sed_targets.reshape(
(sed_targets.shape[0] * sed_targets.shape[1], sed_targets.shape[2])),
sed_outputs.reshape(
(sed_outputs.shape[0] *
sed_outputs.shape[1],
sed_outputs.shape[2])),
average=None)
sed_ap = metrics.average_precision_score(sed_targets.reshape(
(sed_targets.shape[0] * sed_targets.shape[1], sed_targets.shape[2])),
sed_outputs.reshape(
(sed_outputs.shape[0] *
sed_outputs.shape[1],
sed_outputs.shape[2])),
average=None)
logging.info('SED time: {:.3f} s'.format(time.time() - sed_time))
# Evaluate source separation
ss_time = time.time()
(ss_precision, ss_recall, ss_f1_score) = prec_recall_fvalue(
ss_targets.reshape(
(ss_targets.shape[0] * ss_targets.shape[1] * ss_targets.shape[2],
ss_targets.shape[3])),
ss_outputs.reshape(
(ss_outputs.shape[0] * ss_outputs.shape[1] * ss_outputs.shape[2],
ss_outputs.shape[3])),
thres=threshold,
average=None)
logging.info('SS fvalue time: {:.3f} s'.format(time.time() - ss_time))
ss_time = time.time()
ss_auc = metrics.roc_auc_score(
ss_targets.reshape(
(ss_targets.shape[0] * ss_targets.shape[1] * ss_targets.shape[2],
ss_targets.shape[3])),
ss_outputs.reshape(
(ss_outputs.shape[0] * ss_outputs.shape[1] * ss_outputs.shape[2],
ss_outputs.shape[3])),
average=None)
logging.info('SS AUC time: {:.3f} s'.format(time.time() - ss_time))
ss_time = time.time()
ss_ap = metrics.average_precision_score(
ss_targets.reshape(
(ss_targets.shape[0] * ss_targets.shape[1] * ss_targets.shape[2],
ss_targets.shape[3])),
ss_outputs.reshape(
(ss_outputs.shape[0] * ss_outputs.shape[1] * ss_outputs.shape[2],
ss_outputs.shape[3])),
average=None)
logging.info('SS AP time: {:.3f} s'.format(time.time() - ss_time))
# Write stats
stat = {
'at_precision': at_precision,
'at_recall': at_recall,
'at_f1_score': at_f1_score,
'at_auc': at_auc,
'at_ap': at_ap,
'sed_precision': sed_precision,
'sed_recall': sed_recall,
'sed_f1_score': sed_f1_score,
'sed_auc': sed_auc,
'sed_ap': sed_ap,
'ss_precision': ss_precision,
'ss_recall': ss_recall,
'ss_f1_score': ss_f1_score,
'ss_auc': ss_auc,
'ss_ap': ss_ap
}
pickle.dump(stat, open(out_stat_path, 'wb'))
logging.info('Saved stat to {}'.format(out_stat_path))
def train(config, args):
# Wanbd setup
wandb.init(project="covid", entity="sirgarfield", name=args.run_name)
# Arugments & parameters
workspace = config.workspace
model_type = config.model_type
filename = config.filename
holdout_fold = config.holdout_fold
cuda = config.cuda
labels = config.labels
classes_num = len(labels)
seq_len = config.seq_len
mel_bins = config.mel_bins
max_iteration = 50 # To speed up validation, set a maximum iteration
# Paths
hdf5_path = config.h5_path
models_dir = config.model_dir
create_folder(models_dir)
# Model
Model = get_model(model_type)
model = Model(classes_num, seq_len, mel_bins, cuda)
# focal_loss = FocalLoss()
# print('using focal loss')
if config.use_pretrain:
print('Loading pretrained weight ... ')
checkpoint = torch.load(config.pretrain_weight)
pretrained_dict = checkpoint['state_dict']
model_dict = model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if 'final_conv' not in k
}
print(
f' {len(pretrained_dict.keys())}/{len(model_dict.keys())} keys updated '
)
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model.load_state_dict(model_dict)
# wandb.watch(model)
if cuda:
model.cuda()
# Data generator
generator = DataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
holdout_fold=holdout_fold)
# Optimizer
optimizer = optim.Adam(model.parameters(),
lr=1e-3 / 3,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.)
train_bgn_time = time.time()
# Train on mini batches
for iteration, (batch_x, batch_y) in enumerate(generator.generate_train()):
# Evaluate
if iteration % 100 == 0:
train_fin_time = time.time()
(tr_loss, tr_f1_score, tr_auc,
tr_map) = evaluate(model=model,
generator=generator,
data_type='train',
max_iteration=max_iteration,
cuda=cuda)
wandb.log({
'iteration': iteration,
'tr_loss': tr_loss,
'tr_auc': tr_auc,
'tr_f1_score': tr_f1_score,
'tr_map': tr_map
})
logging.info('tr_loss: {:.3f}, tr_f1_score: {:.3f}, '
'tr_auc: {:.3f}, tr_map: {:.3f}'
''.format(tr_loss, tr_f1_score, tr_auc, tr_map))
(va_loss, va_f1_score, va_auc,
va_map) = evaluate(model=model,
generator=generator,
data_type='validate',
max_iteration=max_iteration,
cuda=cuda)
wandb.log({
'iteration': iteration,
'va_loss': va_loss,
'va_auc': va_auc,
'va_f1_score': va_f1_score,
'va_map': va_map
})
logging.info('va_loss: {:.3f}, va_f1_score: {:.3f}, '
'va_auc: {:.3f}, va_map: {:.3f}'
''.format(va_loss, va_f1_score, va_auc, va_map))
train_time = train_fin_time - train_bgn_time
validate_time = time.time() - train_fin_time
logging.info(
'iteration: {}, train time: {:.3f} s, validate time: {:.3f} s'
''.format(iteration, train_time, validate_time))
logging.info('------------------------------------')
train_bgn_time = time.time()
# Save model
if iteration % 1000 == 0 and iteration > 0:
save_out_dict = {
'iteration': iteration,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_out_path = os.path.join(models_dir,
'md_{}_iters.tar'.format(iteration))
torch.save(save_out_dict, save_out_path)
# Save state to wandb
torch.save(
save_out_dict,
os.path.join(wandb.run.dir,
'md_{}_iters.tar'.format(iteration)))
logging.info('Model saved to {}'.format(save_out_path))
# Reduce learning rate
if iteration % 200 == 0 and iteration > 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.9
# Train
batch_x = move_data_to_gpu(batch_x, cuda)
batch_y = move_data_to_gpu(batch_y, cuda)
model.train()
batch_output = model(batch_x)
loss = F.binary_cross_entropy(batch_output, batch_y)
# loss = focal_loss(batch_output, batch_y)
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Stop learning
if iteration == 10000 * 4:
break
wandb.finish()
def train(args):
# Arguments & parameters
workspace = args.workspace
filename = args.filename
validate = args.validate
holdout_fold = args.holdout_fold
cuda = args.cuda
mini_data = args.mini_data
num_classes = len(config.labels)
# Use partial data for training
if mini_data:
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'mini_development.h5')
else:
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'development.h5')
if validate:
validation_csv = os.path.join(workspace, 'validate_meta.csv')
models_dir = os.path.join(workspace, 'models', filename,
'holdout_fold{}'.format(holdout_fold))
else:
validation_csv = None
models_dir = os.path.join(workspace, 'models', filename, 'full_train')
create_folder(models_dir)
# Model
model = Model(num_classes)
if cuda:
model.cuda()
# Optimizer
optimizer = optim.Adam(model.parameters(),
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.)
# Data generator
generator = DataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
time_steps=time_steps,
validation_csv=validation_csv,
holdout_fold=holdout_fold)
iteration = 0
train_bgn_time = time.time()
# Train on mini batches
for (batch_x, batch_y) in generator.generate_train():
# Evaluate
if iteration % 200 == 0:
train_fin_time = time.time()
(tr_acc, tr_mapk) = evaluate(model=model,
generator=generator,
data_type='train',
cuda=cuda)
logging.info('train acc: {:.3f}, train mapk: {:.3f}'.format(
tr_acc, tr_mapk))
if validate:
(va_acc, va_mapk) = evaluate(model=model,
generator=generator,
data_type='validate',
cuda=cuda)
logging.info('valid acc: {:.3f}, validate mapk: {:.3f}'.format(
va_acc, va_mapk))
train_time = train_fin_time - train_bgn_time
validate_time = time.time() - train_fin_time
logging.info('------------------------------------')
logging.info('Iteration: {}, train time: {:.3f} s, eval time: '
'{:.3f} s'.format(iteration, train_time,
validate_time))
train_bgn_time = time.time()
# Save model
if iteration % 1000 == 0 and iteration > 0:
save_out_dict = {
'iteration': iteration,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
save_out_path = os.path.join(models_dir,
'md_{}_iters.tar'.format(iteration))
torch.save(save_out_dict, save_out_path)
logging.info('Save model to {}'.format(save_out_path))
# Reduce learning rate
if iteration % 100 == 0 and iteration > 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.9
batch_x = move_data_to_gpu(batch_x, cuda)
batch_y = move_data_to_gpu(batch_y, cuda)
# Forward
t_forward = time.time()
model.train()
output = model(batch_x)
# Loss
loss = F.nll_loss(output, batch_y)
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
iteration += 1
# Stop learning
if iteration == 10001:
break
def train(args):
# Arugments & parameters
workspace = args.workspace
model_type = args.model_type
filename = args.filename
holdout_fold = args.holdout_fold
scene_type = args.scene_type
snr = args.snr
cuda = args.cuda
labels = config.labels
classes_num = len(labels)
seq_len = config.seq_len
mel_bins = config.mel_bins
max_iteration = 50 # To speed up validation, set a maximum iteration
# Paths
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'scene_type={},snr={}'.format(scene_type, snr),
'development.h5')
models_dir = os.path.join(workspace, 'models', filename,
'model_type={}'.format(model_type),
'scene_type={},snr={}'.format(scene_type, snr),
'holdout_fold{}'.format(holdout_fold))
create_folder(models_dir)
# Model
Model = get_model(model_type)
model = Model(classes_num, seq_len, mel_bins, cuda)
if cuda:
model.cuda()
# Data generator
generator = DataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
holdout_fold=holdout_fold)
# Optimizer
optimizer = optim.Adam(model.parameters(),
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.)
train_bgn_time = time.time()
# Train on mini batches
for iteration, (batch_x, batch_y) in enumerate(generator.generate_train()):
# Evaluate
if iteration % 100 == 0:
train_fin_time = time.time()
(tr_loss, tr_f1_score, tr_auc,
tr_map) = evaluate(model=model,
generator=generator,
data_type='train',
max_iteration=max_iteration,
cuda=cuda)
logging.info('tr_loss: {:.3f}, tr_f1_score: {:.3f}, '
'tr_auc: {:.3f}, tr_map: {:.3f}'
''.format(tr_loss, tr_f1_score, tr_auc, tr_map))
(va_loss, va_f1_score, va_auc,
va_map) = evaluate(model=model,
generator=generator,
data_type='validate',
max_iteration=max_iteration,
cuda=cuda)
logging.info('va_loss: {:.3f}, va_f1_score: {:.3f}, '
'va_auc: {:.3f}, va_map: {:.3f}'
''.format(va_loss, va_f1_score, va_auc, va_map))
train_time = train_fin_time - train_bgn_time
validate_time = time.time() - train_fin_time
logging.info(
'iteration: {}, train time: {:.3f} s, validate time: {:.3f} s'
''.format(iteration, train_time, validate_time))
logging.info('------------------------------------')
train_bgn_time = time.time()
# Save model
if iteration % 1000 == 0 and iteration > 0:
save_out_dict = {
'iteration': iteration,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_out_path = os.path.join(models_dir,
'md_{}_iters.tar'.format(iteration))
torch.save(save_out_dict, save_out_path)
logging.info('Model saved to {}'.format(save_out_path))
# Reduce learning rate
if iteration % 200 == 0 and iteration > 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.9
# Train
batch_x = move_data_to_gpu(batch_x, cuda)
batch_y = move_data_to_gpu(batch_y, cuda)
model.train()
batch_output = model(batch_x)
loss = F.binary_cross_entropy(batch_output, batch_y)
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Stop learning
if iteration == 10000:
break
def forward(model, generate_func, cuda, return_target, return_bottleneck):
"""Forward data to a model.
Args:
model: object
generator_func: function
return_target: bool
return_bottleneck: bool
cuda: bool
Returns:
dict, keys: 'audio_name', 'output'; optional keys: 'target', 'bottleneck'
"""
model.eval()
outputs = []
itemids = []
if return_target:
targets = []
if return_bottleneck:
bottlenecks = []
# Evaluate on mini-batch
for data in generate_func:
if return_target:
(batch_x, batch_y, batch_itemids) = data
else:
(batch_x, batch_itemids) = data
batch_x = move_data_to_gpu(batch_x, cuda)
# Predict
if return_bottleneck:
(batch_output, batch_bottleneck) = model(batch_x,
return_bottleneck=True)
else:
batch_output = model(batch_x, return_bottleneck=False)
outputs.append(batch_output.data.cpu().numpy())
itemids.append(batch_itemids)
if return_target:
targets.append(batch_y)
if return_bottleneck:
bottlenecks.append(batch_bottleneck.data.cpu().numpy())
dict = {}
outputs = np.concatenate(outputs, axis=0)
dict['output'] = outputs
itemids = np.concatenate(itemids, axis=0)
dict['itemid'] = itemids
if return_target:
targets = np.concatenate(targets, axis=0)
dict['target'] = targets
if return_bottleneck:
bottlenecks = np.concatenate(bottlenecks, axis=0)
dict['bottleneck'] = bottlenecks
return dict
def train(args):
# Arugments & parameters
workspace = args.workspace
validate = args.validate
holdout_fold = args.holdout_fold
mini_data = args.mini_data
cuda = args.cuda
filename = args.filename
# Paths
if mini_data:
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'mini_development.h5')
else:
hdf5_path = os.path.join(workspace, 'features', 'logmel',
'development.h5')
if validate:
validation_csv = os.path.join(workspace, 'validation.csv')
models_dir = os.path.join(workspace, 'models', filename,
'holdout_fold{}'.format(holdout_fold))
else:
validation_csv = None
holdout_fold = None
models_dir = os.path.join(workspace, 'models', filename, 'full_train')
create_folder(models_dir)
# Model
model = Model()
if cuda:
model.cuda()
generator = DataGenerator(hdf5_path=hdf5_path,
batch_size=batch_size,
validation_csv=validation_csv,
holdout_fold=holdout_fold)
# Optimizer
optimizer = optim.Adam(model.parameters(),
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.)
# iteration = 0
train_bgn_time = time.time()
# Train on mini-batch
for (iteration, (batch_x,
batch_y)) in enumerate(generator.generate_train()):
# Evaluate
if iteration % 500 == 0:
train_fin_time = time.time()
(tr_acc, tr_auc) = evaluate(model=model,
generator=generator,
data_type='train',
max_iteration=-1,
cuda=cuda)
if validate:
(va_acc, va_auc) = evaluate(model=model,
generator=generator,
data_type='validate',
max_iteration=-1,
cuda=cuda)
train_time = train_fin_time - train_bgn_time
validate_time = time.time() - train_fin_time
# Print info
logging.info(
'iteration: {}, train time: {:.3f} s, validate time: {:.3f} s'
''.format(iteration, train_time, validate_time))
logging.info('tr_acc: {:.3f}, tr_auc: {:.3f}, '.format(
tr_acc, tr_auc))
if validate:
logging.info('va_acc: {:.3f}, va_auc: {:.3f}'.format(
va_acc, va_auc))
logging.info('------------------------------------')
train_bgn_time = time.time()
# Save model
if iteration % 1000 == 0 and iteration > 0:
save_out_dict = {
'iteration': iteration,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_out_path = os.path.join(models_dir,
'md_{}_iters.tar'.format(iteration))
torch.save(save_out_dict, save_out_path)
logging.info('Model saved to {}'.format(save_out_path))
# Reduce learning rate
if iteration % 100 == 0 and iteration > 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.9
# Train
batch_x = move_data_to_gpu(batch_x, cuda)
batch_y = move_data_to_gpu(batch_y, cuda)
model.train()
output = model(batch_x)
loss = F.binary_cross_entropy(output, batch_y)
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Stop learning
if iteration == 10000:
break
