I’m a college student trying to implement Wavenet using PyTorch, this is my first time writing custom modules for a model in PyTorch and I’m having a problem with my model in that it won’t train. Essentially the project is this: I have a set of wav files that I am reading in, processing and quantizing as in the Wavenet paper, and am arranging into series of 1024 data points (the model takes a series of 1024 amplitudes from the wav as input and should output a tensor of 256 values describing the probability that the next item in the series is one of those 256 values).
I’m currently trying to train the model on a single music file, hoping to get it to overfit so that I can be sure that it actually learns from the data. Here lies the problem: the loss won’t decrease as I train. I’ve tried making the model smaller, changing the loss function, changing what kind of layer the output layer is, and making the learning rate larger but nothing seems to work.
I suspect that the problem is somewhere in the model itself, the way that it is constructed may be wrong. It’s possible that I could have linked the custom modules together with the model in a way that interferes with back propagation, at least that is my best guess. My code for the model and my training code is below. I would really appreciate some help!
#model https://github.com/Dankrushen/Wavenet-PyTorch/blob/master/wavenet/models.py
#https://github.com/ryujaehun/wavenet/blob/master/wavenet/networks.py
#https://medium.com/@satyam.kumar.iiitv/understanding-wavenet-architecture-361cc4c2d623
#https://discuss.pytorch.org/t/causal-convolution/3456/4
import torch
import torch.optim as optim
from torch import nn
from functools import reduce
#causal convolution (citation above)
class CausalConv1d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, dilation=1, **kwargs):
super(CausalConv1d, self).__init__()
self.pad = (kernel_size - 1) * dilation
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.dilation = dilation
self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=self.pad, dilation=dilation, **kwargs)
def forward(self, x):
return self.conv(x)
class ResidualBlock(nn.Module):
def __init__(self, input_channels, output_channels, kernel_size, skip_size, skip_channels, dilation=1):
super(ResidualBlock, self).__init__()
self.dilation = dilation
self.skip_size = skip_size
self.conv_s = CausalConv1d(input_channels, output_channels, kernel_size, dilation)#dim
self.sig = nn.Sigmoid()
self.conv_t = CausalConv1d(input_channels, output_channels, kernel_size, dilation)#dim
self.tanh = nn.Tanh()
self.conv_1 = nn.Conv1d(output_channels, output_channels, 1)#dim -> k = 1
self.skip_conv = nn.Conv1d(output_channels, skip_channels, 1)
def forward(self, x):
o = self.sig(self.conv_s(x)) * self.tanh(self.conv_t(x))
skip = self.skip_conv(o)
skip = skip[:,:,-self.skip_size:] #dim control for adding skips
residual = self.conv_1(o)
return residual, skip
class WaveNet(nn.Module):#SET SKIP SIZE default
def __init__(self, skip_size=256, num_blocks=2, num_layers=10, num_hidden=128, kernel_size=2):
super(WaveNet, self).__init__()
self.layer1 = CausalConv1d(1, num_hidden, kernel_size)#dim
self.res_stack = nn.ModuleList()
for b in range(num_blocks):
for i in range(num_layers):
self.res_stack.append(ResidualBlock(num_hidden, num_hidden, kernel_size, skip_size=skip_size, skip_channels=1, dilation=2**i))#dim
#self.hidden = nn.ModuleList(self.hidden)
self.relu1 = nn.ReLU()
self.conv1 = nn.Conv1d(1,1,1)#dim
self.relu2 = nn.ReLU()
self.conv2 = nn.Conv1d(1,1,1)#dim
self.output = nn.Softmax()
def forward(self, x):
skip_vals = []
#initial causal conv
o = self.layer1(x)
#run res blocks
for i, layer in enumerate(self.res_stack):
o, s = layer(o)
skip_vals.append(s)
#sum skip values and pass to last portion of network
o = reduce((lambda a,b: a+b), skip_vals)
o = self.relu1(o)
o = self.conv1(o)
o = self.relu2(o)
o = self.conv2(o)
return self.output(o)
#overfit model to test if it will train
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
net = WaveNet(num_layers=1)
#send to gpu
net.to(device)
criterion = nn.CrossEntropyLoss() #preproc data, stream, train, remember to reformat y label as onehot vector for classification
optimizer = optim.Adam(net.parameters(),lr=0.001)
num_epochs = 20
losses = []
_, inp = wav_to_data(data_path+'/'+wav_files[0])
data = encode(inp)
batch_size = 32
i = 0
for epoch in range(num_epochs):
i += 1
for s in range(0, len(data) - 1024, batch_size):
(x, y) = create_singular_input_stream(data, s, batch_size)
optimizer.zero_grad()
output = net(torch.reshape(torch.FloatTensor(x).to(device), (batch_size,1,1024)))
print(output.shape)
print(torch.Tensor([y]).shape)
#find loss between distributions of amplitudes
#https://discuss.pytorch.org/t/indexerror-target-1-is-out-of-bounds-nlloss/68656
loss = criterion(torch.squeeze(output), torch.squeeze(torch.Tensor(y)).type(torch.LongTensor).to(device))
loss.backward()
optimizer.step()
losses.append(loss.item())
print('Epoch {}/{}, Loss: {:.6f}'.format(i, num_epochs, loss.item()))
A helpful note:
create_singular_input_stream returns x (1,1024 tensor with series) and y (next series value [0,255])
Again, any help is appreciated - I’d love to know what I’m doing wrong.
.