PyTorch WaveNet implementation complications

Anonimouz: <p>Currently, we are a group doing a project about implementing WaveNet in a Tacotron2 → WaveNet → ASR (Given by firm) for midterm project. We are all novices to PyTorch, but recommended to try this library for constructing our WaveNet. We have a problem with the padding and the F.cross_entropy problem for a given .wav-file.</p>
<p>The <strong>main issue</strong> is when we compute the loss function. Our output (from WaveNet) is a tensor of shape:</p>
<pre><code class="lang-auto">output = tensor([1, 256, 225332]) # [batch_size, sample_size, audio_length] 
input = tensor([1, 256, 225360]) 
</code></pre>
<p>There is a problem here, and from what I can see and talk to my supervisor about it is padding the input of the WaveNet. (Cross_entropy wants (N, C) as input and (N) as target, from what I gather, and the dimensions are wrong)<br>
He said “use ‘same’ padding”, but that is currently only operable in TF/Keras as far as I know. We’ve tried to read across multiple posts, but since we’re novices, we can’t seem to figure it out. Any help is appreciated.</p>
<p>This is our WaveNet, which probably has some issues (particularly padding and perhaps causal convolution seems iffy?).</p>
<pre><code class="lang-auto">"""
Wavenet model
"""

from torch import nn
import torch

#TODO: Add local and global conditioning


def initialize(m):
    """
    Initialize CNN with Xavier_uniform weight and 0 bias.
    """
    if isinstance(m, torch.nn.Conv1d):
        nn.init.xavier_uniform_(m.weight)
        nn.init.constant_(m.bias, 0.0)


class CausalConv1d(torch.nn.Module):
    """
    Causal Convolution for WaveNet

    - Jakob
    """

    def __init__(self, in_channels, out_channels, kernel_size, dilation = 1, bias = True):
        super(CausalConv1d, self).__init__()
        # padding=1 for same size(length) between input and output for causal convolution
        self.dilation = dilation
        self.kernel_size = kernel_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.padding = padding = (kernel_size-1) * dilation # kernelsize = 2, -1 * dilation = 1, = 1. - Jakob.
        self.conv = torch.nn.Conv1d(in_channels, out_channels,
                                    kernel_size, padding=padding, dilation=dilation,
                                    bias=bias)  # Fixed for WaveNet but not sure


    def forward(self, x):
        output = self.conv(x)
        if self.padding != 0:
            output = output[:, :, :-self.padding]
        return output





class Wavenet(nn.Module):

    def __init__(self,
                 layers=3,
                 blocks=2,
                 dilation_channels=32,
                 residual_block_channels=512,
                 skip_connection_channels=512,
                 output_channels=256,
                 output_size=32,
                 kernel_size=3
                 ):

        super(Wavenet, self).__init__()

        self.layers = layers
        self.blocks = blocks
        self.dilation_channels = dilation_channels
        self.residual_block_channels = residual_block_channels
        self.skip_connection_channels = skip_connection_channels
        self.output_channels = output_channels
        self.kernel_size = kernel_size
        self.output_size = output_size

        # initialize dilation variables
        receptive_field = 1
        init_dilation = 1


        # List of layers and connections
        self.dilations = []
        self.residual_convs = nn.ModuleList()
        self.filter_conv_layers = nn.ModuleList()
        self.gate_conv_layers = nn.ModuleList()
        self.skip_convs = nn.ModuleList()


        # First convolutional layer
        self.first_conv = CausalConv1d(in_channels=self.output_channels,
                                    out_channels=residual_block_channels,
                                    kernel_size = 2)

        # Building the Modulelists for the residual blocks
        for b in range(blocks):
            additional_scope = kernel_size - 1
            new_dilation = 1
            for i in range(layers):
                # dilations of this layer
                self.dilations.append((new_dilation, init_dilation))

                # dilated convolutions
                self.filter_conv_layers.append(nn.Conv1d(in_channels=residual_block_channels, out_channels=dilation_channels, kernel_size=kernel_size, dilation=new_dilation))

                self.gate_conv_layers.append(nn.Conv1d(in_channels=residual_block_channels, out_channels=dilation_channels, kernel_size=kernel_size, dilation=new_dilation))

                # 1x1 convolution for residual connection
                self.residual_convs.append(nn.Conv1d(in_channels=dilation_channels, out_channels=residual_block_channels, kernel_size=1))

                # 1x1 convolution for skip connection
                self.skip_convs.append(nn.Conv1d(in_channels=dilation_channels,
                                                 out_channels=skip_connection_channels,
                                                 kernel_size=1))

                # Update receptive field and dilation
                receptive_field += additional_scope
                additional_scope *= 2
                init_dilation = new_dilation
                new_dilation *= 2

        # Last two convolutional layers
        self.last_conv_1 = nn.Conv1d(in_channels=skip_connection_channels,
                                  out_channels=skip_connection_channels,
                                  kernel_size=1)

        self.last_conv_2 = nn.Conv1d(in_channels=skip_connection_channels,
                                    out_channels=output_channels,
                                    kernel_size=1)


        #Calculate model receptive field and the required input size for the given output size
        self.receptive_field = receptive_field
        self.input_size = receptive_field + output_size - 1

    def forward(self, input):

        # Feed first convolutional layer with input
        x = self.first_conv(input)

        # Initialize skip connection
        skip = 0

        # Residual block
        for i in range(self.blocks * self.layers):

            (dilation, init_dilation) = self.dilations[i]

            # Residual connection bypassing dilated convolution block
            residual = x

            # input to dilated convolution block
            filter = self.filter_conv_layers[i](x)
            filter = torch.tanh(filter)
            gate = self.gate_conv_layers[i](x)
            gate = torch.sigmoid(gate)
            x = filter * gate

            # Feed into 1x1 convolution for skip connection
            s = self.skip_convs[i](x)

            #Adding skip &amp; Match size with decreasing dimensionality of x
            if skip is not 0:
                skip = skip[:, :, -s.size(2):]
            skip = s + skip # Sum all skip connections

            # Feed into 1x1 convolution for residual connection
            x = self.residual_convs[i](x)
            #Adding Residual &amp; Match size with decreasing dimensionality of x
            x = x + residual[:, :, dilation * (self.kernel_size - 1):]


            # print(x.shape)

        x = torch.relu(skip)
        #Last conv layers
        x = torch.relu(self.last_conv_1(x))
        x = self.last_conv_2(x)
        soft = torch.nn.Softmax(dim=1)
        x = soft(x)
        return x
</code></pre>
<hr>
<p>The training file:</p>
<pre><code class="lang-auto">model = Wavenet(layers=3,blocks=2,output_size=32).to(device)
model.apply(initialize) # xavier_uniform_ : Does this work?
model.train()


optimizer = optim.Adam(model.parameters(), lr=0.0003)
for i, batch in tqdm(enumerate(train_loader)):
    mu_enc_my_x = encode_mu_law(x=batch, mu=256)
    input_tensor = one_hot_encoding(mu_enc_my_x)

    input_tensor = input_tensor.to(device)
    output = model(input_tensor)
    # TODO: Inspect input/output formats, maybe something wrong....
    loss = F.cross_entropy(output.T.reshape(-1, 256), input_tensor[:,:,model.input_size - model.output_size:].long().to(device)) # subtract receptive field instead of pad it, workaround for quick debugging of loss-issue.
    print("\nLoss:", loss.item())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if i % 1000 == 0:
        print("\nSaving model")
        torch.save(model.state_dict(), "wavenet.pt")
</code></pre>

ptrblck: <p>Adding the <code>padding</code> to the conv layers is the right approach, but I’m unsure why you are slicing it in the <code>forward</code>:</p>
<pre><code class="lang-python">    def forward(self, x):
        output = self.conv(x)
        if self.padding != 0:
            output = output[:, :, :-self.padding]
        return output
</code></pre>
<p>This would remove the “right hand side” of the padded output, so is this intended?</p>
<p>Also, <code>nn.CrossEntropyLoss</code> expects raw logits, so remove the <code>nn.Softmax</code> in your model and just pass the output of the last layer to the criterion.</p>

Anonimouz: <p>Yeah the slicing definitely seems fishy. We will look into that, thanks a lot.</p>
<p>Regarding the softmax, it is part of the WaveNet structure as it passes in the output. The architechture below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/2/d2d8cd26742462d28a31dd2bee93e21de5b14396.png" data-download-href="https://discuss.pytorch.org/uploads/default/d2d8cd26742462d28a31dd2bee93e21de5b14396" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d2d8cd26742462d28a31dd2bee93e21de5b14396_2_690x366.png" alt="image" data-base62-sha1="u5eLlddo8NSAaaJAVB04YnSsOWy" width="690" height="366" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d2d8cd26742462d28a31dd2bee93e21de5b14396_2_690x366.png, https://discuss.pytorch.org/uploads/default/original/3X/d/2/d2d8cd26742462d28a31dd2bee93e21de5b14396.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/2/d2d8cd26742462d28a31dd2bee93e21de5b14396.png 2x" data-dominant-color="F7F8F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1013×538 69.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Where should the softmax then be implemented, if output cannot be passed through said function for our loss?<br>
Appreciate the help!!! (Been stuck here for a week)</p>

ptrblck: <p>Yeah, be a bit careful about the architectures presented in papers, as they might not reflect the corresponding framework implementation.<br>
<code>nn.CrossEntropyLoss</code> will apply <code>F.log_softmax</code> and <code>nn.NLLLoss</code> internally, so you should not apply an <code>nn.Softmax</code> layer at the end, since the workflow would then be:</p>
<pre><code class="lang-python">(model -&gt; out -&gt; softmax) -&gt; (log_softmax -&gt; nll_loss)
</code></pre>
<p>where the first part is inside the model and the second in the criterion.<br>
If the mentioned <code>Softmax</code> in the Figure is only used for the output of the model (not internally between layers), then just remove it.</p>

Anonimouz: <p>Alright, that makes sense. We’re essentially double softmaxing. We’ll look into that.</p>
<p>Turning back to the padding problem, I can see what you mean with slicing as mentioned earlier. However, I’ve tried to look a bit at the output when I try to not slice as previously, but just pass the foward as going through the convolutional layer. After running it through the model, I still have the dimensionality issue:</p>
<pre><code class="lang-auto">input = tensor([1, 256, 225360]
output = tensor([1, 256, 225331]) # This is -29, which is exactly our receptive field.
</code></pre>
<p>It evades us still when and how to exactly pad a tensor explicitly. Our intuition says it is before we pass it to the convolutional layer, ie. in my <code>forward(self, x)</code> function:</p>
<pre><code class="lang-auto">class CausalConv1d(torch.nn.Module):
    """
    Causal Convolution for WaveNet
    Causality can be introduced with padding as (kernel_size - 1) * dilation (see Keras documentation)
    or it can be introduced as follows according to Golbin.
    """

    def __init__(self, in_channels, out_channels, kernel_size, dilation = 1, bias = True):
        super(CausalConv1d, self).__init__()
        # padding=1 for same size(length) between input and output for causal convolution
        self.dilation = dilation
        self.kernel_size = kernel_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.padding = (kernel_size-1) * dilation
        self.conv = torch.nn.Conv1d(in_channels, out_channels,
                                    kernel_size, dilation=dilation,
                                    bias=bias)  # Fixed for WaveNet but not sure


    def forward(self, x):
        output = self.conv(x)
        return output
</code></pre>
<p>256 is our “class” parameter, and we know that we can transpose and reshape the tensor to get it in the format <code>nn.torch.functional.Cross_entropy</code> needs. That would be the last brick in the wall.</p>
<p>Yet again, thank you for your patience and help, sir. =)</p>

ptrblck: <p>Usually you would add the <code>padding</code> directly to the initialization of the conv layer.<br>
Unfortunately, there is no option to use the <code>'same'</code> padding argument (there should be some methods to calculate it automatically online for convenience) and you would thus have to calculate the padding manually (as seems to be the case in your <code>CausalConv1d</code>).<br>
However, since you are currently not passing the <code>self.padding</code> to <code>nn.Conv1d</code>, the spatial size of the output might be smaller than the input.</p>