I’m running into another problem with my model during training and generation now and I was wondering whether you could provide any insight on that:
#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
import torch.nn.functional as F
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 , dilation=dilation, **kwargs)
        
    def forward(self, x):
        #pad here to only add to the left side
        x = F.pad(x, (self.pad, 0))
        return self.conv(x)

class ResidualBlock(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size, skip_channels, dilation=1):
        super(ResidualBlock, self).__init__()
        self.dilation = dilation
        self.conv_sig = CausalConv1d(input_channels, output_channels, kernel_size, dilation)#dim
        self.sig = nn.Sigmoid()
        self.conv_tan = CausalConv1d(input_channels, output_channels, kernel_size, dilation)#dim
        self.tanh = nn.Tanh()
        
        #separate weights for residual and skip channels
        self.conv_r = nn.Conv1d(output_channels, output_channels, 1)#dim -> k = 1
        self.conv_s = nn.Conv1d(output_channels, skip_channels, 1)
        
    def forward(self, x):
        o = self.sig(self.conv_sig(x)) * self.tanh(self.conv_tan(x))
        skip = self.conv_s(o)
        #print("SKIP: " + str(skip.shape))
        #skip = skip[:,-self.skip_size:,:] #dim control for adding skips
        #skip = skip[:,:,-1:]
        #print("SK: " + str(skip.shape))
        residual = self.conv_r(o)
        #print("RES: " + str(residual.shape))
        
        return residual, skip
   
    

class WaveNet(nn.Module):
    def __init__(self, skip_channels=256, num_blocks=3, num_layers=10, num_hidden=256, kernel_size=2): 
        super(WaveNet, self).__init__()

        self.embed = nn.Embedding(skip_channels, skip_channels)
        self.layer1 = CausalConv1d(skip_channels, num_hidden, kernel_size)
        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_channels=skip_channels, dilation=2**i))
        
        self.relu1 = nn.ReLU()
        self.conv1 = nn.Conv1d(skip_channels, skip_channels, 1)
        self.relu2 = nn.ReLU()
        self.conv2 = nn.Conv1d(skip_channels, skip_channels, 1)
        #self.output = nn.Softmax()
        
    def forward(self, x):

        o = self.embed(x)
        print('EX: ' + str(o.size()))
        dims = o.size()
        o = o.reshape(dims[0], dims[2], dims[3])
        o = o.transpose(1,2)
        print('ETX: ' + str(o.size()))

        skip_vals = []
        #initial causal conv
        o = self.layer1(o)
        
        #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 o #self.output(o)
    

data_path = '../input/musicwav/80643 Delon  Dalcan - Panik (Original Mix).wav'
model_path = 'model_3B_6L.pt'

#num iterations after which we should save the model
R = 500

#overfit model to test if it will train
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
net = WaveNet(num_blocks=3, num_layers=6)

#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 = 1#00
losses = []

total_steps = 0

if os.path.exists(model_path):
    state = load_model(model_path)
    net.load_state_dict(state['state_dict'])
    optimizer.load_state_dict(state['optimizer'])
    total_steps = state['total_steps']

_, inp = wav_to_data(data_path)
data = encode(inp)
print(data.shape)

batch_size = 64
seq_len = 1024
net.train()

from math import floor
n_steps = floor(len(data) / batch_size)#seq_len)

i = 0
for epoch in range(num_epochs):
    i += 1
    t = 0
    for s in range(0, len(data) - seq_len - batch_size, batch_size):
        t += 1
        #(x, y) = create_singular_input_stream(data, s, batch_size)
        x = []
        y = []
        for g in range(batch_size):
            x.append(data[s+g:s+g+seq_len])
            y.append(data[s+g+1:s+g+seq_len+1])

        #print("X: " + str(torch.Tensor(x).shape))
        optimizer.zero_grad()

        output = net(torch.reshape(torch.LongTensor(x).to(device), (batch_size,1,seq_len)))
        #print(output.shape)
        #print(torch.Tensor(y).shape)

        loss = criterion(output, torch.Tensor(y).type(torch.LongTensor).to(device))

        loss.backward()
        optimizer.step()
        losses.append(loss.item())

        #save model for future training after set amount of iterations
        if t % R == 0:
            save_model(net, optimizer, total_steps+t, model_path)

        print('Epoch {}/{}, Timestep: {}/{}, Loss: {:.6f}'.format(i, num_epochs, t, n_steps, loss.item()))
        



import torch.nn.functional as F

#generate unguided audio
#take first sequence as input and continuously generate a sequence using it
model_path = 'model_3B_6L.pt'
seq_len = 1024
num_samples = 1323000 - seq_len

out_path = 'gen.wav'

#grab first sequence to start with
curr_seq = data[:1024]
generated_seq = curr_seq

net = WaveNet(num_blocks=3, num_layers=6)
state = load_model(model_path)
net.load_state_dict(state['state_dict'])

net.eval()

for i in range(num_samples):
    print('Sample: ' + str(i+1) + ' / ' + str(num_samples))
    output = net(torch.reshape(torch.LongTensor(curr_seq).to(device), (1,1,seq_len)))
    timestep_next = output[0,:,-1:]
    timestep_next = torch.squeeze(timestep_next)
    timestep_next = F.softmax(timestep_next)
    print(timestep_next.size())
    print(timestep_next.cpu().detach())
    t_next = np.argmax(timestep_next.cpu().detach().numpy(), axis=0)
    print(t_next)
    
    #append to sequence
    generated_seq = generated_seq + [t_next]
    
    
#write out file
timeseries_to_wav(np.array(generated_seq), out_path)



When I’m training the model, the loss fluctuates a good deal (though this could be because of the small batch size of 64 I’m using due to memory constraints) but generally does trend downwards. At inference time however, I find that the model produces a constant output of 0. I’ve read around that this could be caused by incorrect use of the quantized inputs, I’ve used the embedding layer rather than passing the input straight to the first convolutional layer as you suggested, but I’m unsure if somewhere I’m making a mistake. One candidate could be the reshaping I am doing in forward() to correct the dimensions of the tensor, another thing I noticed was that you used a TanH activation after your embedding layer which I did not understand - is there a purpose to using an activation function on an embedding vector?