Initial commit

Experiments/audio_regression/configs/bach/val_der_sine.txt

@@ -0,0 +1,5 @@
+exp_name = bach/val_der_sine
+
+supervision = val_der
+filename = gt_bach.wav
+activations = [sine, sine, sine, sine]

Experiments/audio_regression/configs/bach/val_sine.txt

@@ -0,0 +1,5 @@
+exp_name = bach/val_sine
+
+supervision = val
+filename = gt_bach.wav
+activations = [sine, sine, sine, sine]

Experiments/audio_regression/configs/counting/val_der_sine.txt

@@ -0,0 +1,5 @@
+exp_name = counting/val_der_sine
+
+supervision = val_der
+filename = gt_counting.wav
+activations = [sine, sine, sine, sine]

Experiments/audio_regression/configs/counting/val_sine.txt

@@ -0,0 +1,5 @@
+exp_name = counting/val_sine
+
+supervision = val
+filename = gt_counting.wav
+activations = [sine, sine, sine, sine]

Experiments/audio_regression/dataset.py

@@ -0,0 +1,38 @@
+import os
+
+import kornia
+import scipy.io.wavfile as wavfile
+
+import torch
+
+
+def get_data(data_root, filename, factor):
+
+ rate, wav = wavfile.read(os.path.join(data_root, filename)) 
+ print("Rate: %d" % rate)
+ print("Raw data shape: ", wav.shape)
+
+ wav = torch.tensor(wav).reshape(-1, 1)
+ scale = torch.max(torch.abs(wav))
+ wav = wav / scale # (N, 1)
+
+ grad = kornia.filters.spatial_gradient(wav.unsqueeze(0).unsqueeze(0), mode='diff', order=1, normalized=True).squeeze() # (2, N)
+ grad = grad[1, :].reshape(-1, 1) # (N, 1)
+
+ coordinate = torch.linspace(0, len(wav) - 1, len(wav)).reshape(-1, 1) # (N, 1)
+
+ downsampled_wav = wav[::factor, :]
+ downsampled_grad = grad[::factor, :]
+ downsampled_coordinate = coordinate[::factor, :]
+
+ return {
+ 'wav': wav,
+ 'grad': grad,
+ 'coordinate': coordinate,
+
+ 'downsampled_wav': downsampled_wav,
+ 'downsampled_grad': downsampled_grad,
+ 'downsampled_coordinate': downsampled_coordinate,
+ }
+
+

Experiments/audio_regression/diff_operators.py

@@ -0,0 +1,9 @@
+import torch
+
+
+def gradient(y, x, grad_outputs=None):
+ if grad_outputs is None:
+ grad_outputs = torch.ones_like(y)
+ grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]
+ return grad
+

Experiments/audio_regression/loss.py

@@ -0,0 +1,18 @@
+import torch
+
+
+def mse(x, y):
+ return (x - y).pow(2).mean()
+
+
+def val_mse(gt, pred):
+ val_loss = mse(gt, pred)
+
+ return {'val_loss': val_loss}
+
+
+def der_mse(gt_grad, pred_grad):
+ weights = torch.ones(gt_grad.shape[1]).to(gt_grad.device)
+ der_loss = torch.mean((weights * (gt_grad - pred_grad).pow(2)).sum(-1))
+
+ return {'der_loss': der_loss}

Experiments/audio_regression/main.py

@@ -0,0 +1,167 @@
+import os
+import shutil
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+
+from dataset import get_data 
+from model import MLP
+from loss import *
+from utils import *
+
+set_random_seed(0)
+
+
+def config_parser():
+
+ import configargparse
+ parser = configargparse.ArgumentParser()
+ parser.add_argument('--config', is_config_file=True, help="Path of config file.")
+
+ # logging options
+ parser.add_argument('--logging_root', type=str, default='./logs/', help="Where to store ckpts and logs.")
+ parser.add_argument('--epochs_til_ckpt', type=int, default=1000, help="Time interval in epochs until checkpoint is saved.")
+ parser.add_argument('--epochs_til_summary', type=int, default=100, help="Time interval in epochs until tensorboard summary is saved.")
+
+ # training options
+ parser.add_argument('--lrate', type=float, default='5e-5')
+ parser.add_argument('--num_epochs', type=int, default=8000, help="Number of epochs to train for.")
+
+ # experiment options
+ parser.add_argument('--exp_name', type=str, default='supervision_val_der',
+ help="Name of experiment.")
+ parser.add_argument('--supervision', type=str, default='val_der', choices=('val', 'der', 'val_der'))
+ parser.add_argument('--activations', nargs='+', default=['sine', 'sine', 'sine', 'sine'])
+ parser.add_argument('--w0', type=float, default='30.')
+ parser.add_argument('--has_pos_encoding', action='store_true')
+ parser.add_argument('--lambda_der', type=float, default='1.')
+
+ # model options
+ parser.add_argument('--hidden_features', type=int, default=256)
+ parser.add_argument('--num_hidden_layers', type=int, default=3)
+
+ # dataset options
+ parser.add_argument('--data_root', type=str, default='../../data/Audio', help="Root path to audio dataset.")
+ parser.add_argument('--filename', type=str, help="Name of wav file.")
+ parser.add_argument('--factor', type=int, default=5, help="Factor of downsampling.")
+
+ return parser
+
+
+def train(args, model, data, epochs, lrate, epochs_til_summary, epochs_til_checkpoint, logging_dir, train_summary_fn, test_summary_fn, log_f):
+
+ summaries_dir = os.path.join(logging_dir, 'summaries')
+ os.makedirs(summaries_dir)
+ writer = SummaryWriter(summaries_dir)
+
+ checkpoints_dir = os.path.join(logging_dir, 'checkpoints')
+ os.makedirs(checkpoints_dir)
+
+ out_train_imgs_dir = os.path.join(logging_dir, 'out_train_imgs')
+ os.makedirs(out_train_imgs_dir)
+
+ out_test_imgs_dir = os.path.join(logging_dir, 'out_test_imgs')
+ os.makedirs(out_test_imgs_dir)
+
+ optim = torch.optim.Adam(lr=lrate, params=model.parameters())
+
+ # move data to GPU
+ data = {key: value.cuda() for key, value in data.items() if torch.is_tensor(value)}
+
+ for epoch in range(1, epochs + 1):
+
+ # forward and calculate loss
+ model_output = model(data['downsampled_coordinate'], mode='train')
+ losses = {}
+ losses.update(val_mse(data['downsampled_wav'], model_output['pred']))
+ losses.update(der_mse(data['downsampled_grad'], model_output['pred_grad']))
+ if args.supervision == 'val':
+ train_loss = losses['val_loss']
+ elif args.supervision == 'der':
+ train_loss = losses['der_loss']
+ elif args.supervision == 'val_der':
+ train_loss = 1. * losses['val_loss'] + args.lambda_der * losses['der_loss']
+ # tensorboard
+ for loss_name, loss in losses.items():
+ writer.add_scalar(loss_name, loss, epoch)
+ writer.add_scalar("train_loss", train_loss, epoch)
+
+ # backward
+ optim.zero_grad()
+ train_loss.backward()
+ optim.step()
+
+ if (not epoch % epochs_til_summary) or (epoch == epochs):
+
+ # training summary
+ psnr = train_summary_fn(data, model_output, writer, epoch, out_train_imgs_dir)
+ str_print = "[Train] epoch: (%d/%d) " % (epoch, epochs)
+ for loss_name, loss in losses.items():
+ str_print += loss_name + ": %0.6f, " % loss
+ str_print += "PSNR: %.3f " % (psnr)
+ print(str_print)
+ print(str_print, file=log_f)
+
+ # test summary
+ with torch.no_grad():
+ model_output = model(data['coordinate'], mode='test')
+ psnr = test_summary_fn(data, model_output, writer, epoch, out_test_imgs_dir, args.factor, args.filename)
+ str_print = "[Test]: PSNR: %.3f" % (psnr)
+ print(str_print)
+ print(str_print, file=log_f)
+
+ # save checkpoint
+ if (not epoch % epochs_til_checkpoint) or (epoch == epochs):
+ torch.save(model.state_dict(), os.path.join(checkpoints_dir, 'model_epoch_%05d.pth' % epoch))
+
+ torch.save(model.state_dict(), os.path.join(checkpoints_dir, 'model_final.pth'))
+
+
+def main():
+
+ parser = config_parser()
+ args = parser.parse_args()
+
+ logging_dir = os.path.join(args.logging_root, args.exp_name)
+ if os.path.exists(logging_dir):
+ if input("The logging directory %s exists. Overwrite? (y/n)" % logging_dir) == 'y':
+ shutil.rmtree(logging_dir)
+ os.makedirs(logging_dir)
+
+ with open(os.path.join(logging_dir, 'log.txt'), 'w') as log_f:
+
+ print("Args:\n", args)
+ print("Args:\n", args, file=log_f)
+
+ data = get_data(args.data_root, args.filename, args.factor)
+ print('Shape of original wav:', data['wav'].shape)
+ print('Shape of downsampled wav:', data['downsampled_wav'].shape)
+
+ model = MLP(
+ in_features=1,
+ out_features=1,
+ w0=args.w0,
+ activations=args.activations,
+ hidden_features=args.hidden_features,
+ num_hidden_layers=args.num_hidden_layers,
+ has_pos_encoding=args.has_pos_encoding,
+ length=len(data['wav']),
+ fn_samples=len(data['downsampled_wav']))
+ model.cuda()
+
+ train(
+ args=args, 
+ model=model, 
+ data=data, 
+ epochs=args.num_epochs, 
+ lrate=args.lrate, 
+ epochs_til_summary=args.epochs_til_summary, 
+ epochs_til_checkpoint=args.epochs_til_ckpt, 
+ logging_dir=logging_dir, 
+ train_summary_fn=write_train_summary,
+ test_summary_fn=write_test_summary,
+ log_f=log_f)
+
+
+if __name__=='__main__':
+ main()