Vanishing gradient problem

nbravo: <p>Hello everyone,</p>
<p>Currently i’m trying to implement an attention mechanism for video representation but i’m having vanishing gradient issues in the training process, I have been debugging for hours but I still can’t find the problem with my code, that’s why I’m asking for help if it has is any error that could be causing this problem. I’ll explain the mechanism, the paper of reference can be found <a href="https://www-nlpir.nist.gov/projects/tvpubs/tv17.papers/DL-61-86.pdf" rel="noopener nofollow ugc">here</a>.</p>
<p>Let  <img src="https://discuss.pytorch.org/uploads/default/original/3X/0/a/0a5c27de9d00e30b24ada98c9d772b0824591e7a.png" alt="image" data-base62-sha1="1tEdtWIRJdfRNiDPGQlYBoxOyOm" width="68" height="19">  be a set of embedding vectors of a video with N frames. They are extracted using a pretrained convolutional network like Resnet152 or Resnext101. Then, to extract temporal information, they are feed as “word embedding” into a GRU to get a set of N hidden states <img src="https://discuss.pytorch.org/uploads/default/original/3X/0/1/0123437fc51268b96fdd60aa972366bb9bdee650.png" alt="image" data-base62-sha1="a41PwaYwG1ldBMmiWE591ajauY" width="71" height="19">.</p>
<p>With this, they create <img src="https://discuss.pytorch.org/uploads/default/original/3X/1/2/1205490ab0b159785f7d862e5a83a9ec76f341c5.png" alt="image" data-base62-sha1="2zpUVBM0JBDVao4fREnGzjaKJiB" width="169" height="19"> to extract spatial and temporal information of the video.</p>
<p>Then they use an attention mechanism to get a single vector representation for the video. To get the attention map <img src="https://discuss.pytorch.org/uploads/default/original/3X/e/c/ec174b6e1ae0dd1458223d9df8b703e9eee6fcb8.png" alt="image" data-base62-sha1="xGyAG2i4O57aU3IdRsj57T27LIc" width="98" height="22">  they use the mean of the <img src="https://discuss.pytorch.org/uploads/default/original/3X/c/b/cb5ba39d50539dfb216e8bbd6e869ae015ba98ce.png" alt="image" data-base62-sha1="t0ZggHhUKL1BrBuO8H86TV2wh3g" width="17" height="11">  vectors, called <img src="https://discuss.pytorch.org/uploads/default/original/3X/8/6/864adbcc9733b621e173644814e05b4da4b0d4cf.png" alt="image" data-base62-sha1="ja0mCpnPrOCoMRZJS5eDAF3sTtZ" width="9" height="11"> .</p>
<p>They calculate</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/e/aeba7d2ce5f657da6e33b2d8279bf8d483834990.png" alt="image" data-base62-sha1="oVID67MYs2lxmAIATvGOeIKcMlW" width="393" height="18"> <strong>(1)</strong><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/7/b79106dd8daeb12a039096a4b223faa730419225.png" alt="image" data-base62-sha1="qbU5VydaWszyTG1VzXknae95l0F" width="188" height="16">  <strong>(2)</strong><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/7/27182fa96cb6fbcd57da6eac4c6234ad7071eda1.png" alt="image" data-base62-sha1="5zQrhlAW1wR2dUw7Hckz6CDAOpX" width="257" height="56">  <strong>(3)</strong></p>
<p>where <img src="https://discuss.pytorch.org/uploads/default/original/3X/d/4/d496cc2c9c13dd825970d9b1f52026169b5f6000.png" alt="image" data-base62-sha1="ukEiYiklFbSJTxciZfyrUz4Ioj6" width="102" height="17">  are trainable weights and bias, <img src="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd3be88f0a224074a71da24e4d1ef30004c4e885.png" alt="image" data-base62-sha1="r02ATItn20DCIrTBWUO1xrn5P0h" width="25" height="11">  is a trainable vector and <img src="https://discuss.pytorch.org/uploads/default/original/3X/c/5/c5a0c5a9ca8ba9590888af531b21152a5b4fe49c.png" alt="image" data-base62-sha1="scivTcIucKgeD1qO7AF41pNS6cs" width="7" height="7"> is the element-wise product operation.</p>
<p>Then, to get the single vector representation of the set <img src="https://discuss.pytorch.org/uploads/default/original/3X/5/7/572eb8a028cef614312f8ffdc71c9bb5b65b9666.png" alt="image" data-base62-sha1="crfDEIB2ThgBKaD8XSsfyGyypxQ" width="63" height="19"> they calculate:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/f/4f965b75442fd111fbf458258a16500cfac56df4.png" alt="image" data-base62-sha1="bm3RgoPPm7f3sn3WLwdU8otQ7je" width="107" height="54"> <strong>(4)</strong></p>
<p>Now that i explained the mechanism I’m trying to implement, the following code is what I’ve wrote:</p>
<pre><code class="lang-auto">class LinearVector(nn.Module):
    # this module is for the trainable vector W_m
    def __init__(self, size):
        super(LinearVector, self).__init__()
        self.weight = nn.Parameter(nn.init.normal_(torch.rand(size)), requires_grad=True)
        stdw = 1. / math.sqrt(self.weight.size(0))
        self.weight.data.normal_(mean=0, std=stdw)

    def forward(self, m):
        return self.weight * m

class VideoBranch(nn.Module):

    def __init__(self, resnet152_dim, gru_dim, att_dim, out_dim, device):  # 2048, 1024, 1024, 2048, gpu:0
        
        super(VideoBranch, self).__init__()
        self.device = device
        self.gru = nn.GRU(resnet152_dim, gru_dim, batch_first=True)

        self.w_v = nn.Linear(resnet152_dim, att_dim)
        self.w_s = nn.Linear(resnet152_dim + gru_dim, att_dim)
        self.w_m = LinearVector(att_dim)

        self.dropout_1 = nn.Dropout(0.2)
        self.dense_1 = nn.Linear(resnet152_dim + gru_dim, out_dim, bias=False)
        self.dense_1_bn = nn.BatchNorm1d(out_dim)
        self.dropout_2 = nn.Dropout(0.2)
        self.dense_2 = nn.Linear(out_dim, out_dim, bias=False)
        self.dense_2_bn = nn.BatchNorm1d(out_dim)

    def forward(self, resnet152_vectors, n_frames):  #B = batch size, N = max length sequence
        batch_size = resnet152_vectors.size(0)
        N = n_frames.max().item()
        resnet152_vectors = resnet152_vectors[:, :N, :] # delete unnecesary padding.  size = [B, N, 2048]
        embs = pack_padded_sequence(resnet152_vectors, n_frames, batch_first=True, enforce_sorted=False)
        h_packed, _ = self.gru(embs)
        h, _ = pad_packed_sequence(h_packed, batch_first=True) # size of h = [B, N, 1024]

        # attention mechanism
        s = torch.cat((h, resnet152_vectors), 2).view(batch_size, N, -1) # size = [B, N, 1024 + 2048]
        v = resnet152_vectors.sum(dim=1) / n_frames.view(batch_size, -1).float()  # resnet152_vectors.mean() ignoring padding. size = [B, 2048]
        y_s = self.w_s(s).tanh()  # size = [B, N, 1024] 1024 is the att_size
        y_v = self.w_v(v).tanh()  # size = [B, 1024]
        m = torch.mul(y_s, y_v.view(batch_size, 1, -1)) # equation (1), size=[B, N, 1024]
     
       # from here, the gradient vanishes  for W_s, W_v and the GRU
     
        r = self.w_m(m)  # equation (2), size=[B, N, 1024]
        alpha = torch.sum(r, dim=2).softmax(dim=1) # equation (3), size = [B, N]
        alpha = alpha.view(batch_size, N, 1)   # size = [B, N, 1]
        v_a = torch.sum(torch.mul(alpha, s), dim=1) # equation (4), size = [B, 2048]

        # fully-connected mapping to another vector space
        y = self.dense_1(v_a)
        y = self.dense_1_bn(y)
        y = nn.ReLU()(y)
        y = self.dropout_1(y)
        y = self.dense_2(y)
        y = self.dense_2_bn(y)
        y = nn.ReLU()(y)
        y = self.dropout_2(y)
        return y
</code></pre>
<p>Using a gradient flow plot shared by <a href="https://discuss.pytorch.org/t/check-gradient-flow-in-network/15063/7">RoshanRane</a>, at the start of the training process the gradient starts to vanish  for <img src="https://discuss.pytorch.org/uploads/default/original/3X/a/e/aed2c13b5817435d5f8d16284e127f9bf05ee5f7.png" alt="image" data-base62-sha1="oWyCr7NlujjH9esG2Suq3Z1jjW7" width="56" height="16">  and the GRU.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/7/f7da737aeae24ecdabb887aa1d199c2c27bf7d81.png" alt="image" data-base62-sha1="zmBXNAd07FKT82JrrvdWG0rBJ7j" width="315" height="475"></p>
<p>and after a few batches the gradient completely vanish<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/5/d5e22d372ba9e34f3b058f1724b3fa1b8553b9c0.png" alt="image" data-base62-sha1="uw6hvuqEJI3mUTq2CICIIcXRDUc" width="315" height="475"></p>
<p>Is there any error in this implementation? Help would be very appreciated.<br>
Thank you very much!</p>