mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-18 07:33:20 +00:00
first commit; melgan support (perhaps bloated)
I copied this from my old and messy `tacotron-melgan` repo, where I had previously experimented with running the two models in a pipeline. I'll be moving over all the JIT pieces here and cleaning up when I can.
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,3 @@
|
||||
*~
|
||||
__pycache__/
|
||||
python/
|
||||
@@ -0,0 +1,7 @@
|
||||
jit-models
|
||||
==========
|
||||
|
||||
This project exists to turn PyTorch models into Torch JIT, which can be run in
|
||||
any language with libtorch bindings.
|
||||
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self):
|
||||
super(Discriminator, self).__init__()
|
||||
|
||||
self.discriminator = nn.ModuleList([
|
||||
nn.Sequential(
|
||||
nn.ReflectionPad1d(7),
|
||||
nn.utils.weight_norm(nn.Conv1d(1, 16, kernel_size=15, stride=1)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.utils.weight_norm(nn.Conv1d(16, 64, kernel_size=41, stride=4, padding=20, groups=4)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.utils.weight_norm(nn.Conv1d(64, 256, kernel_size=41, stride=4, padding=20, groups=16)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.utils.weight_norm(nn.Conv1d(256, 1024, kernel_size=41, stride=4, padding=20, groups=64)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.utils.weight_norm(nn.Conv1d(1024, 1024, kernel_size=41, stride=4, padding=20, groups=256)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.utils.weight_norm(nn.Conv1d(1024, 1024, kernel_size=5, stride=1, padding=2)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
),
|
||||
nn.utils.weight_norm(nn.Conv1d(1024, 1, kernel_size=3, stride=1, padding=1)),
|
||||
])
|
||||
|
||||
def forward(self, x):
|
||||
'''
|
||||
returns: (list of 6 features, discriminator score)
|
||||
we directly predict score without last sigmoid function
|
||||
since we're using Least Squares GAN (https://arxiv.org/abs/1611.04076)
|
||||
'''
|
||||
features = list()
|
||||
for module in self.discriminator:
|
||||
x = module(x)
|
||||
features.append(x)
|
||||
return features[:-1], features[-1]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = Discriminator()
|
||||
|
||||
x = torch.randn(3, 1, 22050)
|
||||
print(x.shape)
|
||||
|
||||
features, score = model(x)
|
||||
for feat in features:
|
||||
print(feat.shape)
|
||||
print(score.shape)
|
||||
|
||||
pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
print(pytorch_total_params)
|
||||
@@ -0,0 +1,108 @@
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .res_stack import ResStack
|
||||
|
||||
import numpy as np
|
||||
|
||||
from wav_images import render_histogram
|
||||
from wav_images import rescale_mel
|
||||
|
||||
def load_melgan(model_checkpoint_path):
|
||||
print('Loading melgan checkpoint: {}'.format(model_checkpoint_path))
|
||||
|
||||
checkpoint = torch.load(model_checkpoint_path, map_location=torch.device('cpu'))
|
||||
melgan_state = checkpoint['model_g']
|
||||
melgan_state = transform_melgate_state(melgan_state)
|
||||
|
||||
model_melgan = Generator()
|
||||
model_melgan.load_state_dict(melgan_state)
|
||||
return model_melgan
|
||||
|
||||
def transform_melgate_state(melgan_state):
|
||||
import re
|
||||
BLOCK_REGEX = re.compile('blocks\.(\d+)\.')
|
||||
SHORTCUT_REGEX = re.compile('shortcuts\.(\d+)\.')
|
||||
|
||||
new_model_g = [] # Rebuild the ordered dict
|
||||
for key, value in melgan_state.items():
|
||||
if 'blocks.' in key:
|
||||
new_key = BLOCK_REGEX.sub(r'block\1.', key)
|
||||
#print('{} -> {}'.format(key, new_key))
|
||||
key = new_key
|
||||
elif 'shortcuts.' in key:
|
||||
new_key = SHORTCUT_REGEX.sub(r'shortcut\1.', key)
|
||||
#print('{} -> {}'.format(key, new_key))
|
||||
key = new_key
|
||||
|
||||
new_model_g.append((key, value))
|
||||
|
||||
return OrderedDict(new_model_g)
|
||||
|
||||
class Generator(nn.Module):
|
||||
def __init__(self):
|
||||
super(Generator, self).__init__()
|
||||
mel_channel = 80 # TODO: This is an hparam
|
||||
self.mel_channel = mel_channel
|
||||
|
||||
self.generator = nn.Sequential(
|
||||
nn.ReflectionPad1d(3),
|
||||
nn.utils.weight_norm(nn.Conv1d(mel_channel, 512, kernel_size=7, stride=1)),
|
||||
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.utils.weight_norm(nn.ConvTranspose1d(512, 256, kernel_size=16, stride=8, padding=4)),
|
||||
|
||||
ResStack(256),
|
||||
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.utils.weight_norm(nn.ConvTranspose1d(256, 128, kernel_size=16, stride=8, padding=4)),
|
||||
|
||||
ResStack(128),
|
||||
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.utils.weight_norm(nn.ConvTranspose1d(128, 64, kernel_size=4, stride=2, padding=1)),
|
||||
|
||||
ResStack(64),
|
||||
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.utils.weight_norm(nn.ConvTranspose1d(64, 32, kernel_size=4, stride=2, padding=1)),
|
||||
|
||||
ResStack(32),
|
||||
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.ReflectionPad1d(3),
|
||||
nn.utils.weight_norm(nn.Conv1d(32, 1, kernel_size=7, stride=1)),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def forward(self, mel):
|
||||
mel = (mel + 5.0) / 5.0 # roughly normalize spectrogram
|
||||
audio = self.generator(mel)
|
||||
audio = audio.squeeze()
|
||||
|
||||
MAX_WAV_VALUE = 32768.0
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
|
||||
max_ = torch.max(torch.abs(audio))
|
||||
#max_ = audio.abs().max()
|
||||
#max_ = max(0.01, max_)
|
||||
multiplier = 32767.0 / max_
|
||||
audio = audio * multiplier
|
||||
#max_ = np.max(np.abs(audio))
|
||||
#max_ = max(0.01, audio)
|
||||
#print("Max: {}".format(max_))
|
||||
#print("Multiplier: {}".format(multiplier))
|
||||
|
||||
return audio
|
||||
|
||||
def eval(self):
|
||||
super(Generator, self).eval()
|
||||
|
||||
def inference(self, mel):
|
||||
audio = self.forward(mel)
|
||||
#audio = audio.squeeze() # collapse all dimension except time axis
|
||||
return audio
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Identity(nn.Module):
|
||||
def __init__(self):
|
||||
super(Identity, self).__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .discriminator import Discriminator
|
||||
from .identity import Identity
|
||||
|
||||
|
||||
class MultiScaleDiscriminator(nn.Module):
|
||||
def __init__(self):
|
||||
super(MultiScaleDiscriminator, self).__init__()
|
||||
|
||||
self.discriminators = nn.ModuleList(
|
||||
[Discriminator() for _ in range(3)]
|
||||
)
|
||||
|
||||
self.pooling = nn.ModuleList(
|
||||
[Identity()] +
|
||||
[nn.AvgPool1d(kernel_size=4, stride=2, padding=1, count_include_pad=False) for _ in range(1, 3)]
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
ret = list()
|
||||
|
||||
for pool, disc in zip(self.pooling, self.discriminators):
|
||||
x = pool(x)
|
||||
ret.append(disc(x))
|
||||
|
||||
return ret # [(feat, score), (feat, score), (feat, score)]
|
||||
@@ -0,0 +1,39 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
class ResStack(nn.Module):
|
||||
def __init__(self, channel):
|
||||
super(ResStack, self).__init__()
|
||||
|
||||
block_list = [
|
||||
nn.Sequential(
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.ReflectionPad1d(3**i),
|
||||
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=3, dilation=3**i)),
|
||||
nn.LeakyReLU(0.2),
|
||||
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1)),
|
||||
)
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
self.block0 = block_list[0]
|
||||
self.block1 = block_list[1]
|
||||
self.block2 = block_list[2]
|
||||
|
||||
shortcut_list = [
|
||||
nn.utils.weight_norm(nn.Conv1d(channel, channel, kernel_size=1))
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
self.shortcut0 = shortcut_list[0]
|
||||
self.shortcut1 = shortcut_list[1]
|
||||
self.shortcut2 = shortcut_list[2]
|
||||
|
||||
def forward(self, x):
|
||||
x = self.shortcut0(x) + self.block0(x)
|
||||
x = self.shortcut1(x) + self.block1(x)
|
||||
x = self.shortcut2(x) + self.block2(x)
|
||||
return x
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
# Requirements from Tacotron
|
||||
Unidecode==0.4.20
|
||||
#falcon==1.2.0
|
||||
falcon==2.0.0
|
||||
inflect==0.2.5
|
||||
librosa==0.5.1
|
||||
numpy==1.14.3
|
||||
tensorflow==1.14.0
|
||||
|
||||
# Requirements from MelGan
|
||||
PyYAML==5.2
|
||||
torch==1.3.1
|
||||
tqdm==4.40.2
|
||||
|
||||
# Personal Extensions
|
||||
Pillow==6.2.1
|
||||
@@ -0,0 +1,181 @@
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
def render_histogram(wav, filename='histograms/output.png', row_index=0, col_index=1):
|
||||
dimensions = wav.shape
|
||||
print('>>> Image input size: {}'.format(dimensions))
|
||||
rows = dimensions[row_index]
|
||||
cols = dimensions[col_index]
|
||||
|
||||
# Remove 1-sized dimensions
|
||||
wav_signal = wav.squeeze()
|
||||
|
||||
print('>>> Generating image (rows: {} cols: {})'.format(rows, cols))
|
||||
image = Image.new('RGB', (rows, cols))
|
||||
pixels = image.load()
|
||||
|
||||
minimum = wav_signal[0,0]
|
||||
maximum = wav_signal[0,0]
|
||||
|
||||
for x in range(0, rows):
|
||||
for y in range(0, cols):
|
||||
if wav_signal[x,y] > maximum:
|
||||
maximum = wav_signal[x,y]
|
||||
if wav_signal[x,y] < minimum:
|
||||
minimum = wav_signal[x,y]
|
||||
|
||||
print('Minimum: ' + str(minimum))
|
||||
print('Maximum: ' + str(maximum))
|
||||
|
||||
if maximum == minimum:
|
||||
maximum = 1 # nb: to prevent division by zero below
|
||||
|
||||
for x in range(0, rows):
|
||||
for y in range(0, cols):
|
||||
v = wav_signal[x,y]
|
||||
scaled = int((v - minimum) / (maximum - minimum) * 255)
|
||||
pixels[x, y] = (scaled, scaled, scaled)
|
||||
|
||||
# Image Show only works on local X server:
|
||||
#image.show()
|
||||
image.save(filename)
|
||||
|
||||
def rescale_mel(mel, rescaled_min=-1.0, rescaled_max=1.0):
|
||||
# NB: Assumes 3 dimensions
|
||||
dimensions = mel.shape
|
||||
minimum = mel[0,0,0]
|
||||
maximum = mel[0,0,0]
|
||||
|
||||
for x in range(0, dimensions[0]):
|
||||
for y in range(0, dimensions[1]):
|
||||
for z in range(0, dimensions[2]):
|
||||
if mel[x,y,z] > maximum:
|
||||
maximum = mel[x,y,z]
|
||||
if mel[x,y,z] < minimum:
|
||||
minimum = mel[x,y,z]
|
||||
|
||||
print('Actual Minimum: ' + str(minimum))
|
||||
print('Actual Maximum: ' + str(maximum))
|
||||
|
||||
old_range = float(maximum) - minimum
|
||||
rescaled_range = float(rescaled_max) - rescaled_min
|
||||
|
||||
print('Old Range: ' + str(old_range))
|
||||
print('Rescaled Range: ' + str(rescaled_range))
|
||||
|
||||
for x in range(0, dimensions[0]):
|
||||
for y in range(0, dimensions[1]):
|
||||
for z in range(0, dimensions[2]):
|
||||
d = mel[x,y,z]
|
||||
rescaled = (d - minimum) * rescaled_range / old_range + rescaled_min
|
||||
mel[x,y,z] = rescaled
|
||||
|
||||
"""
|
||||
# The following is taken from SciPy Cookbook
|
||||
# https://scipy-cookbook.readthedocs.io/items/Rebinning.html
|
||||
import numpy as n
|
||||
import scipy.interpolate
|
||||
import scipy.ndimage
|
||||
|
||||
def congrid(a, newdims, method='linear', centre=False, minusone=False):
|
||||
'''Arbitrary resampling of source array to new dimension sizes.
|
||||
Currently only supports maintaining the same number of dimensions.
|
||||
To use 1-D arrays, first promote them to shape (x,1).
|
||||
|
||||
Uses the same parameters and creates the same co-ordinate lookup points
|
||||
as IDL''s congrid routine, which apparently originally came from a VAX/VMS
|
||||
routine of the same name.
|
||||
|
||||
method:
|
||||
neighbour - closest value from original data
|
||||
nearest and linear - uses n x 1-D interpolations using
|
||||
scipy.interpolate.interp1d
|
||||
(see Numerical Recipes for validity of use of n 1-D interpolations)
|
||||
spline - uses ndimage.map_coordinates
|
||||
|
||||
centre:
|
||||
True - interpolation points are at the centres of the bins
|
||||
False - points are at the front edge of the bin
|
||||
|
||||
minusone:
|
||||
For example- inarray.shape = (i,j) & new dimensions = (x,y)
|
||||
False - inarray is resampled by factors of (i/x) * (j/y)
|
||||
True - inarray is resampled by(i-1)/(x-1) * (j-1)/(y-1)
|
||||
This prevents extrapolation one element beyond bounds of input array.
|
||||
'''
|
||||
if not a.dtype in [n.float64, n.float32]:
|
||||
a = n.cast[float](a)
|
||||
|
||||
m1 = n.cast[int](minusone)
|
||||
ofs = n.cast[int](centre) * 0.5
|
||||
old = n.array( a.shape )
|
||||
ndims = len( a.shape )
|
||||
if len( newdims ) != ndims:
|
||||
print("[congrid] dimensions error. " \
|
||||
"This routine currently only support " \
|
||||
"rebinning to the same number of dimensions.")
|
||||
return None
|
||||
newdims = n.asarray( newdims, dtype=float )
|
||||
dimlist = []
|
||||
|
||||
if method == 'neighbour':
|
||||
for i in range( ndims ):
|
||||
base = n.indices(newdims)[i]
|
||||
dimlist.append( (old[i] - m1) / (newdims[i] - m1) \
|
||||
* (base + ofs) - ofs )
|
||||
cd = n.array( dimlist ).round().astype(int)
|
||||
newa = a[list( cd )]
|
||||
return newa
|
||||
|
||||
elif method in ['nearest','linear']:
|
||||
# calculate new dims
|
||||
for i in range( ndims ):
|
||||
base = n.arange( newdims[i] )
|
||||
dimlist.append( (old[i] - m1) / (newdims[i] - m1) \
|
||||
* (base + ofs) - ofs )
|
||||
# specify old dims
|
||||
olddims = [n.arange(i, dtype = n.float) for i in list( a.shape )]
|
||||
|
||||
# first interpolation - for ndims = any
|
||||
mint = scipy.interpolate.interp1d( olddims[-1], a, kind=method )
|
||||
newa = mint( dimlist[-1] )
|
||||
|
||||
trorder = [ndims - 1] + list(range( ndims - 1 ))
|
||||
for i in range( ndims - 2, -1, -1 ):
|
||||
newa = newa.transpose( trorder )
|
||||
|
||||
mint = scipy.interpolate.interp1d( olddims[i], newa, kind=method )
|
||||
newa = mint( dimlist[i] )
|
||||
|
||||
if ndims > 1:
|
||||
# need one more transpose to return to original dimensions
|
||||
newa = newa.transpose( trorder )
|
||||
|
||||
return newa
|
||||
elif method in ['spline']:
|
||||
oslices = [ slice(0,j) for j in old ]
|
||||
oldcoords = n.ogrid[oslices]
|
||||
nslices = [ slice(0,j) for j in list(newdims) ]
|
||||
newcoords = n.mgrid[nslices]
|
||||
|
||||
newcoords_dims = list(range(n.rank(newcoords)))
|
||||
#make first index last
|
||||
newcoords_dims.append(newcoords_dims.pop(0))
|
||||
newcoords_tr = newcoords.transpose(newcoords_dims)
|
||||
# makes a view that affects newcoords
|
||||
|
||||
newcoords_tr += ofs
|
||||
|
||||
deltas = (n.asarray(old) - m1) / (newdims - m1)
|
||||
newcoords_tr *= deltas
|
||||
|
||||
newcoords_tr -= ofs
|
||||
|
||||
newa = scipy.ndimage.map_coordinates(a, newcoords)
|
||||
return newa
|
||||
else:
|
||||
print("Congrid error: Unrecognized interpolation type.\n", \
|
||||
"Currently only \'neighbour\', \'nearest\',\'linear\',", \
|
||||
"and \'spline\' are supported.")
|
||||
return None
|
||||
"""
|
||||
|
||||
Executable
+43
@@ -0,0 +1,43 @@
|
||||
#!/usr/bin/env python3
|
||||
# From: https://github.com/pytorch/pytorch/issues/20356#issuecomment-545572400
|
||||
import torch
|
||||
|
||||
from collections import OrderedDict
|
||||
import re
|
||||
import yaml
|
||||
|
||||
from model_melgan.generator import load_melgan
|
||||
|
||||
#torch.set_default_tensor_type('torch.DoubleTensor')
|
||||
|
||||
#model_checkpoint_path = '/home/bt/models/melgan-swpark/trumptxlearn183_fa14347_7675.pt'
|
||||
model_checkpoint_path = '/home/bt/models/melgan-swpark/ljs-high-pitch_8741385_6975.pt'
|
||||
model_checkpoint_path = '/home/bt/models/melgan-swpark/trump-txlearn-2020.05.05_8741385_7525.pt'
|
||||
model_checkpoint_path = '/home/bt/models/melgan-swpark/trump-txlearn-2020.05.05_8741385_13675.pt'
|
||||
model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_trump_txlearn_2020-05-30_steno_chkpt9800.pt'
|
||||
|
||||
melgan_model = load_melgan(model_checkpoint_path)
|
||||
|
||||
#output_filename = 'shared_melgan_container2.pt'
|
||||
#output_filename = 'melgan_ljs-high-pitch_container.pt'
|
||||
output_filename = 'melgan_trump-txlearn-2020.05.05.jit'
|
||||
output_filename = 'melgan_trump-txlearn-2020.05.05_13675.jit'
|
||||
output_filename = 'melgan_trump_txlearn_2020-05-30_steno_chkpt9800.torchjit'
|
||||
|
||||
# NB: Tracing evaluates the model on input and unrolls and hardcodes branching
|
||||
# and loops. Scripting allows these to remain by converting the entire program
|
||||
# to TorchScript. Scripting seems MUCH faster, is also harder get working.
|
||||
trace_model = True
|
||||
if trace_model:
|
||||
print('Tracing model and saving as: {}'.format(output_filename))
|
||||
mel_file = '/home/bt/dev/voder/data/mels/LJ002-0320.mel'
|
||||
mel = torch.load(mel_file, map_location=torch.device('cpu'))
|
||||
# NB: Getting `Tracing failed sanity checks! Graphs differed across invocations!`
|
||||
traced_script_module = torch.jit.trace(melgan_model, mel, check_trace=True)
|
||||
traced_script_module.save(output_filename)
|
||||
else:
|
||||
print('Saving script model as: {}'.format(output_filename))
|
||||
container = torch.jit.script(melgan_model)
|
||||
container.save(output_filename)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user