Add main folder for repo
This commit is contained in:
parent
0a54d23b22
commit
af26da8cad
16 changed files with 809 additions and 0 deletions
1
main/CSR_Net/__init__.py
Normal file
1
main/CSR_Net/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .map import MlRes
|
||||
BIN
main/CSR_Net/__init__.pyc
Normal file
BIN
main/CSR_Net/__init__.pyc
Normal file
Binary file not shown.
BIN
main/CSR_Net/__pycache__/__init__.cpython-37.pyc
Normal file
BIN
main/CSR_Net/__pycache__/__init__.cpython-37.pyc
Normal file
Binary file not shown.
BIN
main/CSR_Net/__pycache__/blocks.cpython-37.pyc
Normal file
BIN
main/CSR_Net/__pycache__/blocks.cpython-37.pyc
Normal file
Binary file not shown.
BIN
main/CSR_Net/__pycache__/map.cpython-37.pyc
Normal file
BIN
main/CSR_Net/__pycache__/map.cpython-37.pyc
Normal file
Binary file not shown.
BIN
main/CSR_Net/__pycache__/util.cpython-37.pyc
Normal file
BIN
main/CSR_Net/__pycache__/util.cpython-37.pyc
Normal file
Binary file not shown.
81
main/CSR_Net/blocks.py
Normal file
81
main/CSR_Net/blocks.py
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
import tensorflow as tf
|
||||
from tensorflow.keras import layers, initializers
|
||||
|
||||
from CSR_Net.util import SubPixel1D
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Encoder(layers.Layer):
|
||||
'''encodes input audio to higher-dimensions'''
|
||||
|
||||
def __init__(self, n_filters, n_filtersizes, name='encoder', **kwargs):
|
||||
super(Encoder, self).__init__(name=name, **kwargs)
|
||||
self.w = layers.Conv1D(filters=n_filters, kernel_size=n_filtersizes, padding='same',
|
||||
kernel_initializer=initializers.Orthogonal(gain=1.0, seed=None), strides=2)
|
||||
self.a = layers.LeakyReLU(0.2)
|
||||
|
||||
def call(self, inputs):
|
||||
x = self.w(inputs)
|
||||
return self.a(x)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Bottleneck(layers.Layer):
|
||||
'''middle layer of the network, used to sample data mostly'''
|
||||
|
||||
def __init__(self, n_filters, n_filtersizes, name='bottleneck', **kwargs):
|
||||
super(Bottleneck, self).__init__(name=name, **kwargs)
|
||||
self.w = layers.Conv1D(filters=n_filters, kernel_size=n_filtersizes, padding='same',
|
||||
kernel_initializer=initializers.Orthogonal(gain=1.0, seed=None), strides=2)
|
||||
self.n = layers.Dropout(0.5)
|
||||
self.a = layers.LeakyReLU(0.2)
|
||||
|
||||
def call(self, inputs):
|
||||
x = self.w(inputs)
|
||||
x = self.n(x)
|
||||
return self.a(x)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
class Decoder(layers.Layer):
|
||||
'''decodes (upsamples) high-dimension data back down to audio dimension'''
|
||||
|
||||
def __init__(self, n_filters, n_filtersizes, name='decoder', **kwargs):
|
||||
super(Decoder, self).__init__(name=name, **kwargs)
|
||||
self.w = layers.Conv1D(filters=2*n_filters, kernel_size=n_filtersizes, padding='same',
|
||||
kernel_initializer=initializers.Orthogonal(gain=1.0, seed=None))
|
||||
self.n = layers.Dropout(0.5)
|
||||
self.a = layers.Activation('relu')
|
||||
# (-1, n, f)
|
||||
self.s = layers.Lambda(SubPixel1D, arguments={'r':2})
|
||||
|
||||
def call(self, inputs):
|
||||
x = self.w(inputs)
|
||||
x = self.n(x)
|
||||
x = self.a(x)
|
||||
return self.s(x)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
class OutputConv(layers.Layer):
|
||||
'''output layer for the network'''
|
||||
|
||||
def __init__(self, n_filters, n_filtersizes, name='finalconv', **kwargs):
|
||||
super(OutputConv, self).__init__(name=name, **kwargs)
|
||||
self.w = layers.Conv1D(filters=2, kernel_size=9, padding='same',
|
||||
kernel_initializer=initializers.Orthogonal(gain=1.0, seed=None))
|
||||
self.s = layers.Lambda(SubPixel1D, arguments={'r':2})
|
||||
|
||||
def call(self, inputs):
|
||||
x = self.w(inputs)
|
||||
return self.s(x)
|
||||
75
main/CSR_Net/map.py
Normal file
75
main/CSR_Net/map.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
import tensorflow as tf
|
||||
import numpy as np
|
||||
|
||||
from CSR_Net import util, blocks
|
||||
|
||||
|
||||
class MlRes(tf.keras.Model):
|
||||
'''the model: piecing together the blocks and defining the logic'''
|
||||
|
||||
def __init__(self):
|
||||
super(MlRes, self).__init__()
|
||||
|
||||
# utils
|
||||
self.merge1 = util.MergeTensors(type='concat')
|
||||
self.merge2 = util.MergeTensors(type='add')
|
||||
|
||||
# n_kernels = [ 64, 128, 256, 384, 384, 384, 384, 384]
|
||||
# n_filters = [ 128, 256, 512, 512, 512, 512, 512, 512]
|
||||
# n_filters = [ 256, 512, 512, 512, 512, 1024, 1024, 1024]
|
||||
# n_filtersizes = [129, 65, 33, 17, 9, 9, 9, 9]
|
||||
# n_filtersizes = [31, 31, 31, 31, 31, 31, 31, 31]
|
||||
# kernel_size = [65, 33, 17, 9, 9, 9, 9, 9, 9]
|
||||
|
||||
# blocks
|
||||
self.encode1 = blocks.Encoder(128, 65) # (num_filters, filter_size)
|
||||
self.encode2 = blocks.Encoder(256, 33)
|
||||
self.encode3 = blocks.Encoder(512, 17)
|
||||
# self.encode4 = blocks.Encoder(512, 9)
|
||||
|
||||
self.bottleneck = blocks.Bottleneck(512, 9)
|
||||
|
||||
# self.decode4 = blocks.Decoder(512, 9)
|
||||
self.decode3 = blocks.Decoder(512, 17)
|
||||
self.decode2 = blocks.Decoder(256, 33)
|
||||
self.decode1 = blocks.Decoder(128, 65)
|
||||
|
||||
self.finalconv = blocks.OutputConv(2, 9)
|
||||
|
||||
def call(self, inputs):
|
||||
skip = []
|
||||
|
||||
x = self.encode1(inputs)
|
||||
skip.append(x)
|
||||
|
||||
x = self.encode2(x)
|
||||
skip.append(x)
|
||||
|
||||
x = self.encode3(x)
|
||||
skip.append(x)
|
||||
|
||||
# x = self.encode4(x)
|
||||
# skip.append(x)
|
||||
|
||||
|
||||
x = self.bottleneck(x)
|
||||
|
||||
|
||||
# x = self.decode4(x)
|
||||
# x = self.merge1([x, skip[-1]])
|
||||
|
||||
x = self.decode3(x)
|
||||
x = self.merge1([x, skip[-1]])
|
||||
|
||||
x = self.decode2(x)
|
||||
x = self.merge1([x, skip[-2]])
|
||||
|
||||
x = self.decode1(x)
|
||||
x = self.merge1([x, skip[-3]])
|
||||
|
||||
|
||||
x = self.finalconv(x)
|
||||
x = self.merge2([x, inputs])
|
||||
|
||||
|
||||
return x
|
||||
BIN
main/CSR_Net/map.pyc
Normal file
BIN
main/CSR_Net/map.pyc
Normal file
Binary file not shown.
34
main/CSR_Net/util.py
Normal file
34
main/CSR_Net/util.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
import tensorflow as tf
|
||||
from tensorflow.keras import layers
|
||||
|
||||
def SubPixel1D(I, r):
|
||||
'''
|
||||
One-dimensional subpixel upsampling layer
|
||||
Calls a tensorflow function that directly implements this functionality.
|
||||
We assume input has dim (batch, width, r)
|
||||
'''
|
||||
X = tf.transpose(I, [2,1,0]) # (r, w, b)
|
||||
X = tf.batch_to_space(X, [r], [[0,0]]) # (1, r*w, b)
|
||||
X = tf.transpose(X, [2,1,0])
|
||||
return X
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras import layers
|
||||
|
||||
|
||||
class MergeTensors(layers.Layer):
|
||||
"""custom layer that handles merging tensors for upscaling purposes"""
|
||||
def __init__(self, type, name='merger', **kwargs):
|
||||
super(MergeTensors, self).__init__(name=name, **kwargs)
|
||||
self.type = type
|
||||
self.c = layers.Concatenate(axis=-1)
|
||||
self.a = layers.Add()
|
||||
|
||||
def call(self, inputs):
|
||||
if self.type == 'concat':
|
||||
x = self.c(inputs)
|
||||
if self.type == 'add':
|
||||
x = self.a(inputs)
|
||||
return x
|
||||
BIN
main/__pycache__/ds.cpython-37.pyc
Normal file
BIN
main/__pycache__/ds.cpython-37.pyc
Normal file
Binary file not shown.
BIN
main/__pycache__/prep_vctk.cpython-37.pyc
Normal file
BIN
main/__pycache__/prep_vctk.cpython-37.pyc
Normal file
Binary file not shown.
BIN
main/__pycache__/util.cpython-37.pyc
Normal file
BIN
main/__pycache__/util.cpython-37.pyc
Normal file
Binary file not shown.
160
main/ds.py
Normal file
160
main/ds.py
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
import os, argparse
|
||||
import numpy as np
|
||||
import h5py
|
||||
import random
|
||||
|
||||
import librosa
|
||||
from scipy import interpolate
|
||||
from scipy.signal import decimate
|
||||
|
||||
from scipy.signal import butter, lfilter
|
||||
|
||||
|
||||
class Prep_VCTK:
|
||||
"""main class for creating data pipelines"""
|
||||
# we just changed the add_data function into the __init__ method
|
||||
def __init__(self, type, num_files, dim, file_list,
|
||||
scale=4,
|
||||
interpolate=True,
|
||||
low_pass=False,
|
||||
batch_size=32,
|
||||
sr=16000,
|
||||
sam=0.25):
|
||||
self.type = type
|
||||
self.num_files = num_files
|
||||
self.scale = scale
|
||||
self.dim = dim
|
||||
self.stride = dim
|
||||
self.interpolate = interpolate
|
||||
self.low_pass = low_pass
|
||||
self.batch_size = batch_size
|
||||
self.sr = sr
|
||||
self.sam = sam
|
||||
|
||||
self.path = '../data/multispeaker'
|
||||
out = f'{self.path}/vctk-{self.type}.{self.scale}.{self.sr}.{self.dim}.{self.num_files}.{self.sam}.h5'
|
||||
with h5py.File(out, 'w') as f: # create h5 file for data to be placed in
|
||||
self.add_data(h5_file=f, inputfiles=file_list, save_examples=False)
|
||||
|
||||
def add_data(self, h5_file, inputfiles, save_examples=False):
|
||||
# Make a list of all files to be processed
|
||||
file_list = []
|
||||
file_extensions = set(['.wav'])
|
||||
with open(inputfiles) as f:
|
||||
for line in f: # for every file in the .txt file
|
||||
filename = line.strip() # strips any spaces off of the filename
|
||||
ext = os.path.splitext(filename)[1]
|
||||
if ext in file_extensions: # if file is wavefile, add to file_list
|
||||
file_list.append(filename) # add path of filename before adding
|
||||
file_list = random.sample(file_list, int(self.num_files))
|
||||
|
||||
# patches to extract and their size
|
||||
if self.interpolate: # if user wants to replace low-res patches with cubpic splines
|
||||
d, d_lr = self.dim, self.dim # dimensions for lr and sd are the same
|
||||
s, s_lr = self.stride, self.stride # extracting low-res stride
|
||||
else: # apply scaling to lr audio
|
||||
d, d_lr = self.dim, self.dim / self.scale
|
||||
s, s_lr = self.stride, self.stride / self.scale
|
||||
hr_patches, lr_patches = list(), list()
|
||||
|
||||
for j, file_path in enumerate(file_list): #update user on progress (ie. 30/240)
|
||||
if j % 10 == 0:
|
||||
print (f'Making {self.type} data...{int(np.ceil(j/self.num_files*100))}% \r', end='')
|
||||
# load audio file from file_list
|
||||
x, fs = librosa.load(f'../data/{file_path}', sr=self.sr)
|
||||
|
||||
# crop so that it works with scaling ratio (ie. divisible by 2, 4, 6, etc.)
|
||||
x_len = len(x) # length of file
|
||||
x = x[ : x_len - (x_len % self.scale)]
|
||||
|
||||
# generate low-res version
|
||||
if self.low_pass:
|
||||
# x_bp = butter_bandpass_filter(x, 0, args.sr / args.scale / 2, fs, order=6)
|
||||
# x_lr = np.array(x[0::args.scale])
|
||||
#x_lr = decimate(x, args.scale, zero_phase=True)
|
||||
x_lr = decimate(x, self.scale) # downsample signal after applying anti-aliasing filter
|
||||
else:
|
||||
x_lr = np.array(x[0::self.scale]) #just sample audio at a lower rate (every 2, 4, 6, etc.)
|
||||
|
||||
if self.interpolate: # zero padd array to have same dim as HD (ie, 4000300020001)
|
||||
x_lr = Prep_VCTK.upsample(self, x_lr)
|
||||
assert len(x) % self.scale == 0
|
||||
assert len(x_lr) == len(x)
|
||||
else:
|
||||
assert len(x) % self.scale == 0
|
||||
assert len(x_lr) == len(x) / self.scale
|
||||
|
||||
# generate patches
|
||||
max_i = len(x) - int(d) + 1 # max iteration?: file length - dimension + 1
|
||||
# iterate through the file in strides
|
||||
for i in range(0, max_i, s):
|
||||
# keep only a fraction of all the patches (not in use)
|
||||
u = np.random.uniform() # a single value is returned between 0 and 1
|
||||
if u > self.sam: continue # only keeping a random % of the patches if args.sam is specified
|
||||
|
||||
if self.interpolate:
|
||||
i_lr = i
|
||||
else:
|
||||
i_lr = i / self.scale
|
||||
|
||||
hr_patch = np.array( x[i : i+d] ) # current patch = current position + dim
|
||||
lr_patch = np.array( x_lr[i_lr : i_lr+d_lr] )
|
||||
|
||||
# print 'a', hr_patch
|
||||
# print 'b', lr_patch
|
||||
|
||||
assert len(hr_patch) == d
|
||||
assert len(lr_patch) == d_lr
|
||||
|
||||
# print hr_patch
|
||||
|
||||
hr_patches.append(hr_patch.reshape((d,1))) # create hr patches
|
||||
lr_patches.append(lr_patch.reshape((d_lr,1))) # create lr patches
|
||||
|
||||
# if j == 1: exit(1)
|
||||
|
||||
# crop # of patches so that it's a multiple of mini-batch size
|
||||
num_patches = len(hr_patches)
|
||||
print (f'num_patches = {num_patches}')
|
||||
num_to_keep = int(np.floor(num_patches / self.batch_size) * self.batch_size)
|
||||
hr_patches = np.array(hr_patches[:num_to_keep])
|
||||
lr_patches = np.array(lr_patches[:num_to_keep])
|
||||
|
||||
print (hr_patches.shape)
|
||||
|
||||
# create the hdf5 file
|
||||
data_set = h5_file.create_dataset('data', lr_patches.shape, np.float32)
|
||||
label_set = h5_file.create_dataset('label', hr_patches.shape, np.float32)
|
||||
|
||||
# fill hdf5 files with patches
|
||||
data_set[...] = lr_patches
|
||||
label_set[...] = hr_patches
|
||||
|
||||
def upsample(self, x_lr): #lr = lowres, hr = highres
|
||||
x_lr = x_lr.flatten() # flatten audio array
|
||||
x_hr_len = len(x_lr) * self.scale # get (len of audio array * scaling factor)
|
||||
x_sp = np.zeros(x_hr_len) # create zero-padded array with new length
|
||||
|
||||
i_lr = np.arange(x_hr_len, step=self.scale) # create lr array with step size of scaling factor
|
||||
i_hr = np.arange(x_hr_len)
|
||||
|
||||
f = interpolate.splrep(i_lr, x_lr) # "Given the set of data points (x[i], y[i]) determine a smooth spline approximation"
|
||||
|
||||
# Given the knots and coefficients of a B-spline representation, evaluate the value of the smoothing polynomial and its derivatives.
|
||||
x_sp = interpolate.splev(i_hr, f)
|
||||
|
||||
return x_sp
|
||||
|
||||
@staticmethod
|
||||
def butter_bandpass(lowcut, highcut, fs, order=5):
|
||||
nyq = 0.5 * fs
|
||||
low = lowcut / nyq
|
||||
high = highcut / nyq
|
||||
b, a = butter(order, [low, high], btype='band')
|
||||
return b, a
|
||||
|
||||
@staticmethod
|
||||
def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
|
||||
b, a = butter_bandpass(lowcut, highcut, fs, order=order)
|
||||
y = lfilter(b, a, data)
|
||||
return y
|
||||
141
main/main.py
Normal file
141
main/main.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
import os
|
||||
import random
|
||||
import datetime
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras import Model
|
||||
from tensorboard.plugins.hparams import api as hp
|
||||
|
||||
import CSR_Net
|
||||
import util
|
||||
|
||||
# confirm tf is using GPU
|
||||
# print("Num GPUs Available: ", len(tf.config.experimental.list_physical_device$
|
||||
# input('Press enter of gpu settings are good')
|
||||
|
||||
# gundersena@75.86.178.105:~/Desktop/crimata-super-res/train/logs/weights ~/Desktop
|
||||
# scp rm -r gundersena.75.86.178.105:~/Desktop/crimata-super-res/main
|
||||
# scp -r ~/Desktop/crimata-super-res/main gundersena@75.86.178.105:~/Desktop/crimata-super-res
|
||||
|
||||
|
||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
|
||||
|
||||
def make_parser():
|
||||
"""creates argument parser from train and eval"""
|
||||
parser = argparse.ArgumentParser()
|
||||
subparsers = parser.add_subparsers(title='Commands')
|
||||
|
||||
# train
|
||||
train_parser = subparsers.add_parser('train')
|
||||
train_parser.set_defaults(func=train)
|
||||
|
||||
train_parser.add_argument('-i','--model-id')
|
||||
train_parser.add_argument('-c','--from_ckpt')
|
||||
train_parser.add_argument('-k','--new-data')
|
||||
train_parser.add_argument('-d','--dim-size',type=int)
|
||||
train_parser.add_argument('-x','--num-files',type=int)
|
||||
# train_parser.add_argument('-t','--train-file')
|
||||
# train_parser.add_argument('-v','--val-file')
|
||||
train_parser.add_argument('-e','--epochs',type=int)
|
||||
train_parser.add_argument('-b','--batch-size',type=int)
|
||||
train_parser.add_argument('-o','--cycle-length',type=int)
|
||||
train_parser.add_argument('-m','--max-lr',type=float)
|
||||
train_parser.add_argument('-n','--min-lr',type=float)
|
||||
|
||||
# eval
|
||||
eval_parser = subparsers.add_parser('eval')
|
||||
eval_parser.set_defaults(func=eval)
|
||||
|
||||
eval_parser.add_argument('-i','--model-id')
|
||||
eval_parser.add_argument('-n','--num-examples',type=int)
|
||||
eval_parser.add_argument('-w','--wavfile-list')
|
||||
eval_parser.add_argument('-r','--scale',type=int)
|
||||
eval_parser.add_argument('-s','--sample-rate',type=int)
|
||||
eval_parser.add_argument('-a','--make-audio')
|
||||
eval_parser.add_argument('-c','--from-ckpt', default='True')
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def train(args):
|
||||
"""High-level method for training a model"""
|
||||
# load data
|
||||
x_train, y_train, n_sam = util.load_data(args, type='train', num_files=args.num_files, full_data=True)
|
||||
x_val, y_val = util.load_data(args, type='val', num_files=int(np.floor(args.num_files*0.3)))
|
||||
|
||||
# callbacks
|
||||
checkpointer = tf.keras.callbacks.ModelCheckpoint(filepath=f'logs/weights/weights.{args.model_id}.tf',
|
||||
monitor='val_loss', save_best_only=True, save_weights_only=True, mode='auto')
|
||||
|
||||
# smart_learn = CSR_Net.util.SGDRScheduler(min_lr=args.min_lr, max_lr=args.max_lr,
|
||||
# steps_per_epoch=np.ceil(n_sam/args.batch_size), cycle_length=args.cycle_length)
|
||||
|
||||
# lr_finder = CSR_Net.util.LRFinder(min_lr=1e-7, max_lr=3e-2,
|
||||
# steps_per_epoch=np.ceil(n_sam/args.batch_size), epochs=args.epochs)
|
||||
|
||||
# logdir = f'logs/fit/{args.model_id}-{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}'
|
||||
# hparams = {'max_lr':args.max_lr, 'min_lr':args.min_lr, 'cycle_length':args.cycle_length}
|
||||
# param_logger = hp.KerasCallback(logdir, hparams)
|
||||
#
|
||||
# tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir, histogram_freq=1,
|
||||
# write_graph=True, update_freq='epoch')
|
||||
|
||||
# make model
|
||||
model = make_model(args)
|
||||
|
||||
# compile model
|
||||
optimizer = tf.keras.optimizers.Adam(learning_rate=args.max_lr)
|
||||
# optimizer = tf.keras.optimizers.SGD(learning_rate=args.max_lr, momentum=0.8, nesterov=False)
|
||||
model.compile(optimizer=optimizer, loss='mean_squared_error')
|
||||
|
||||
# final review
|
||||
util.review_model(args, x_train, y_train)
|
||||
|
||||
# train model
|
||||
model.fit(x=x_train, y=y_train, batch_size=args.batch_size, epochs=args.epochs,
|
||||
callbacks=[checkpointer],
|
||||
validation_data=[x_val, y_val], shuffle=True)
|
||||
|
||||
# plot loss and lr metrics
|
||||
# lr_finder.plot_lr()
|
||||
# lr_finder.plot_loss()
|
||||
|
||||
|
||||
def eval(args):
|
||||
"""test the model on real audio"""
|
||||
# make model
|
||||
model = make_model(args)
|
||||
|
||||
# create list of file names
|
||||
file_list = []
|
||||
with open(args.wavfile_list) as f:
|
||||
for line in f:
|
||||
file_list.append(line) # this is gonna get pretty big for a real dataset...
|
||||
|
||||
# eval on random sample of files
|
||||
file_list = random.sample(file_list, args.num_examples)
|
||||
for idx, line in enumerate(file_list):
|
||||
file = line.rstrip('\n')
|
||||
CSR_Net.util.eval_wav(file, args, model)
|
||||
|
||||
|
||||
def make_model(args):
|
||||
"""define a graph and compile model"""
|
||||
model = CSR_Net.MlRes()
|
||||
|
||||
if args.from_ckpt == 'True':
|
||||
model.load_weights((f'logs/weights/weights.{args.model_id}.tf'))
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def main():
|
||||
parser = make_parser()
|
||||
args = parser.parse_args()
|
||||
args.func(args)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
317
main/util.py
Normal file
317
main/util.py
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
from keras.callbacks import Callback
|
||||
import keras.backend as K
|
||||
import numpy as np
|
||||
|
||||
|
||||
class SGDRScheduler(Callback):
|
||||
"""custom callback for implementing a SGDR learning rate"""
|
||||
def __init__(self, min_lr, max_lr, steps_per_epoch, lr_decay=0.9, cycle_length=10,
|
||||
mult_factor=1.5):
|
||||
self.min_lr = min_lr
|
||||
self.max_lr = max_lr
|
||||
self.lr_decay = lr_decay
|
||||
|
||||
self.batch_since_restart = 0
|
||||
self.next_restart = cycle_length
|
||||
|
||||
self.steps_per_epoch = steps_per_epoch
|
||||
|
||||
self.cycle_lenrfrrgth = cycle_length
|
||||
self.mult_factor = mult_factor
|
||||
|
||||
|
||||
def clr(self):
|
||||
fraction_to_restart = self.batch_since_restart / (self.steps_per_epoch * self.cycle_length)
|
||||
lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + np.cos(fraction_to_restart * np.pi))
|
||||
return lr
|
||||
|
||||
|
||||
def on_train_begin(self, logs=None):
|
||||
K.set_value(self.model.optimizer.lr, self.max_lr)
|
||||
|
||||
|
||||
def on_batch_end(self, batch, logs=None):
|
||||
self.batch_since_restart += 1
|
||||
K.set_value(self.model.optimizer.lr, self.clr())
|
||||
|
||||
|
||||
def on_epoch_end(self, epoch, logs=None):
|
||||
if epoch + 1 == self.next_restart:
|
||||
self.batch_since_restart = 0
|
||||
self.cycle_length = np.ceil(self.cycle_length * self.mult_factor)
|
||||
self.next_restart += self.cycle_length
|
||||
self.max_lr *= self.lr_decay
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
import matplotlib.pyplot as plt
|
||||
import keras.backend as K
|
||||
from keras.callbacks import Callback
|
||||
|
||||
|
||||
class LRFinder(Callback):
|
||||
"""
|
||||
custom callback for evaluating the optimal lr range for SGDR
|
||||
Usage:
|
||||
lr_finder = models.util.LRFinder(min_lr=1e-5, max_lr=3e-2,
|
||||
steps_per_epoch=np.ceil(n_sam/args.batch_size), epochs=3)
|
||||
"""
|
||||
def __init__(self, min_lr=1e-5, max_lr=1e-2, steps_per_epoch=None, epochs=None):
|
||||
super(LRFinder, self).__init__()
|
||||
|
||||
self.min_lr = min_lr
|
||||
self.max_lr = max_lr
|
||||
self.total_iterations = steps_per_epoch * epochs
|
||||
self.iteration = 0
|
||||
self.history = {}
|
||||
|
||||
|
||||
def clr(self):
|
||||
'''Calculate the learning rate.'''
|
||||
x = self.iteration / self.total_iterations
|
||||
return self.min_lr + (self.max_lr-self.min_lr) * x
|
||||
|
||||
|
||||
def on_train_begin(self, logs=None):
|
||||
'''Initialize the learning rate to the minimum value at the start of training.'''
|
||||
logs = logs or {}
|
||||
K.set_value(self.model.optimizer.lr, self.min_lr)
|
||||
|
||||
|
||||
def on_batch_end(self, epoch, logs=None):
|
||||
'''Record previous batch statistics and update the learning rate.'''
|
||||
logs = logs or {}
|
||||
self.iteration += 1
|
||||
|
||||
self.history.setdefault('lr', []).append(K.get_value(self.model.optimizer.lr))
|
||||
self.history.setdefault('iterations', []).append(self.iteration)
|
||||
|
||||
for k, v in logs.items():
|
||||
self.history.setdefault(k, []).append(v)
|
||||
|
||||
K.set_value(self.model.optimizer.lr, self.clr())
|
||||
|
||||
|
||||
def plot_lr(self):
|
||||
'''Helper function to quickly inspect the learning rate schedule.'''
|
||||
plt.plot(self.history['iterations'], self.history['lr'])
|
||||
plt.yscale('log')
|
||||
plt.xlabel('Iteration')
|
||||
plt.ylabel('Learning rate')
|
||||
plt.tight_layout()
|
||||
plt.savefig('plots/lr.png')
|
||||
plt.clf()
|
||||
|
||||
def plot_loss(self):
|
||||
'''Helper function to quickly observe the learning rate experiment results.'''
|
||||
plt.plot(self.history['lr'], self.history['loss'])
|
||||
plt.xscale('log')
|
||||
plt.xlabel('Learning rate')
|
||||
plt.ylabel('Loss')
|
||||
plt.tight_layout()
|
||||
plt.savefig('plots/loss.png')
|
||||
plt.clf()
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import h5py
|
||||
import ds
|
||||
|
||||
def load_data(args, type, num_files, full_data=False):
|
||||
np.set_printoptions(threshold=100)
|
||||
path = '../data/multispeaker'
|
||||
|
||||
# load training data
|
||||
datasets = os.listdir(path)
|
||||
|
||||
for dataset in datasets:
|
||||
if str(args.dim_size) and str(num_files) in dataset:
|
||||
if args.new_data == 'False':
|
||||
make_data = False
|
||||
break
|
||||
else:
|
||||
make_data = True
|
||||
|
||||
if make_data:
|
||||
ds.Prep_VCTK(type=type, num_files=num_files, dim=args.dim_size, file_list=f'{path}/{type}-files.txt')
|
||||
|
||||
with h5py.File(f'{path}/vctk-{type}.4.16000.{args.dim_size}.{num_files}.0.25.h5', 'r') as hf:
|
||||
X = np.array(hf.get('data'))
|
||||
Y = np.array(hf.get('label'))
|
||||
|
||||
n_sam, n_dim, n_chan = Y.shape
|
||||
r = Y[0].shape[1] / X[0].shape[1]
|
||||
|
||||
if full_data:
|
||||
return X, Y, n_sam
|
||||
else:
|
||||
return X, Y
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
import os
|
||||
|
||||
|
||||
def review_model(args, x_train, y_train):
|
||||
"""reviews model parameters and raises warnings if something is not recommended"""
|
||||
# prints preivew of the data
|
||||
preview_data(x_train, y_train)
|
||||
|
||||
# assert not overwriting weights
|
||||
if args.from_ckpt == 'False':
|
||||
files = os.listdir('./logs/weights')
|
||||
for file in files:
|
||||
if f'loss.{args.model_id}' in file:
|
||||
input('Warning: Are you sure you want to write over these weights?')
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
import numpy as np
|
||||
|
||||
|
||||
def preview_data(X, Y):
|
||||
print ('Preview X:')
|
||||
print (f'Shape: {X.shape}')
|
||||
print (f'Max: {np.amax(X)} | Min: {np.amin(X)}')
|
||||
print (X[1])
|
||||
|
||||
print ('Preview Y:')
|
||||
print (f'Shape of Y: {Y.shape}')
|
||||
print (f'Max: {np.amax(Y)} | Min: {np.amin(Y)}')
|
||||
print (Y[1])
|
||||
|
||||
# data = eval_wav.get_spectrum(X[:100].flatten(), n_fft=2048)
|
||||
# label = eval_wav.get_spectrum(Y[:100].flatten(), n_fft=2048)
|
||||
|
||||
input('Press enter to continue...')
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
import os
|
||||
import librosa
|
||||
import numpy as np
|
||||
from keras.models import Model
|
||||
from scipy import interpolate
|
||||
from scipy.signal import decimate
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
|
||||
class eval_wav:
|
||||
'''
|
||||
Helper function for eval() in main.py
|
||||
Takes a single wavfile and evaluates it by exporting audio and spectrogram
|
||||
for hr, lr, and pr
|
||||
'''
|
||||
def __init__(self, file, args, model):
|
||||
# ../data/VCTK-Corpus---
|
||||
x_hr, fs = librosa.load(file, sr=args.sample_rate)
|
||||
|
||||
# ensure that input is a multiple of 2^downsampling layers
|
||||
ds_layers = 5
|
||||
x_hr = eval_wav.clip(x_hr, 2**ds_layers)
|
||||
assert len(x_hr) % 2**ds_layers == 0
|
||||
|
||||
# downscale signal
|
||||
# x_lr = decimate(x_hr, args.scale)
|
||||
x_lr = np.array(x_hr[0::args.scale])
|
||||
# x_lr = downsample_bt(x_hr, args.scale)
|
||||
assert len(x_hr)/len(x_lr) == args.scale
|
||||
|
||||
|
||||
# upsample signal through interpolation
|
||||
x_ir = eval_wav.upsample(x_lr, args.scale)
|
||||
assert len(x_ir) == len(x_hr)
|
||||
|
||||
# trim array again to make it a multiple of 800
|
||||
x_ir = eval_wav.clip(x_ir, 800)
|
||||
print(f'Input length: {len(x_ir)}')
|
||||
|
||||
|
||||
n_sam = len(x_ir)/800
|
||||
x_pr = model.predict(x_ir.reshape(int(n_sam), 800, 1))
|
||||
x_pr = x_pr.flatten()
|
||||
|
||||
# save the file
|
||||
filename = os.path.basename(file)
|
||||
name = os.path.splitext(filename)[-2]
|
||||
|
||||
if args.make_audio:
|
||||
audio_data = np.concatenate((x_hr, x_ir, x_pr), axis=0)
|
||||
audio_outname = f'../samples/audio/{name}'
|
||||
librosa.output.write_wav(audio_outname + '.hr.wav', audio_data, fs)
|
||||
|
||||
# save the spectrum
|
||||
spec_outname = f'../samples/spectrograms/{name}'
|
||||
self.outfile=spec_outname + '.png'
|
||||
self.S_pr = eval_wav.get_spectrum(x_pr, n_fft=2048)
|
||||
self.S_hr = eval_wav.get_spectrum(x_hr, n_fft=2048)
|
||||
self.S_lr = eval_wav.get_spectrum(x_lr, n_fft=2048/args.scale)
|
||||
self.S_ir = eval_wav.get_spectrum(x_ir, n_fft=2048)
|
||||
self.save_spectrum()
|
||||
|
||||
@staticmethod
|
||||
def upsample(x_lr, r): #lr = lowres, hr = highres
|
||||
|
||||
x_lr = x_lr.flatten() # flatten audio array
|
||||
x_hr_len = len(x_lr) * r # get (len of audio array * scaling factor)
|
||||
x_sp = np.zeros(x_hr_len) # create zero-padded array with new length
|
||||
|
||||
i_lr = np.arange(x_hr_len, step=r) # create lr array with step size of scaling factor
|
||||
i_hr = np.arange(x_hr_len)
|
||||
|
||||
f = interpolate.splrep(i_lr, x_lr) # "Given the set of data points (x[i], y[i]) determine a smooth spline approximation"
|
||||
|
||||
# Given the knots and coefficients of a B-spline representation, evaluate the value of the smoothing polynomial and its derivatives.
|
||||
x_sp = interpolate.splev(i_hr, f)
|
||||
|
||||
return x_sp
|
||||
|
||||
@staticmethod
|
||||
def clip(array, multiple):
|
||||
x_len = len(array)
|
||||
remainder = x_len % multiple
|
||||
x_len = x_len - remainder
|
||||
array = array[:x_len]
|
||||
return array
|
||||
|
||||
@staticmethod
|
||||
def get_spectrum(data, n_fft=2048):
|
||||
S = librosa.stft(data, int(n_fft))
|
||||
S = np.log1p(np.abs(S))
|
||||
p = np.angle(S)
|
||||
S = np.log1p(np.abs(S))
|
||||
return S.T
|
||||
|
||||
def save_spectrum(self, lim=1000):
|
||||
plt.subplot(2,2,1)
|
||||
plt.title('Target')
|
||||
plt.xlabel('Frequency')
|
||||
plt.ylabel('Time')
|
||||
plt.imshow(self.S_hr, aspect=10)
|
||||
plt.xlim([0,lim])
|
||||
|
||||
plt.subplot(2,2,2)
|
||||
plt.title('Test')
|
||||
plt.xlabel('Frequency')
|
||||
plt.ylabel('Time')
|
||||
plt.imshow(self.S_lr, aspect=10)
|
||||
plt.xlim([0,lim])
|
||||
|
||||
plt.subplot(2,2,3)
|
||||
plt.title('Interp')
|
||||
plt.xlabel('Frequency')
|
||||
plt.ylabel('Time')
|
||||
plt.imshow(self.S_ir, aspect=10)
|
||||
plt.xlim([0,lim])
|
||||
|
||||
plt.subplot(2,2,4)
|
||||
plt.title('Predict')
|
||||
plt.xlabel('Frequency')
|
||||
plt.ylabel('Time')
|
||||
plt.imshow(self.S_pr, aspect=10)
|
||||
plt.xlim([0,lim])
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(self.outfile)
|
||||
Loading…
Reference in a new issue