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 here.
Let 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 .
With this, they create to extract spatial and temporal information of the video.
Then they use an attention mechanism to get a single vector representation for the video. To get the attention map they use the mean of the vectors, called .
where are trainable weights and bias, is a trainable vector and is the element-wise product operation.
Then, to get the single vector representation of the set they calculate:
Now that i explained the mechanism I’m trying to implement, the following code is what I’ve wrote:
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
Using a gradient flow plot shared by RoshanRane, at the start of the training process the gradient starts to vanish for and the GRU.
and after a few batches the gradient completely vanish
Is there any error in this implementation? Help would be very appreciated.
Thank you very much!