I have followed your idea but the loss is the same quite weird.

loss_index_mask.PNG1446×577 50.4 KB


  
      github.com
  
  
    pcshih/pytorch-VSLUD/blob/20a91393c040cd6a005c1b01ce77ca3db82efd25/SK.py#L24-L56

h = x
x_temp = x


h = self.FCSN(h)


values, indices = h.max(1, keepdim=True)


###old###
# # 0/1 vector, we only want key(indices=1) frame
# column_mask = (indices==1).view(-1).nonzero().view(-1).tolist() 


# # if S_K doesn't select more than one element, then random select two element(for the sake of diversity loss)
# if len(column_mask)<2:
#     print("S_K does not select anything, give a random mask with 2 elements")
#     column_mask = random.sample(list(range(h.shape[3])), 2)


# index = torch.tensor(column_mask, device=torch.device('cuda:0'))
# h_select = torch.index_select(h, 3, index)
# x_select = torch.index_select(x_temp, 3, index)
###old###


  This file has been truncated. show original

  
  
    
    
  
  



  
      github.com
  
  
    pcshih/pytorch-VSLUD/blob/20a91393c040cd6a005c1b01ce77ca3db82efd25/train.py#L202-L214

###new reconstruct###
reconstruct_loss = torch.sum((S_K_summary-vd)**2 * index_mask) / torch.sum(index_mask)
###new reconstruct###




# diversity
S_K_summary_reshape = S_K_summary.view(S_K_summary.shape[1], S_K_summary.shape[3])
norm_div = torch.norm(S_K_summary_reshape, 2, 0, True)
S_K_summary_reshape = S_K_summary_reshape/norm_div
loss_matrix = S_K_summary_reshape.transpose(1, 0).mm(S_K_summary_reshape)
diversity_loss = loss_matrix.sum() - loss_matrix.trace()
#diversity_loss = diversity_loss/len(column_mask)/(len(column_mask)-1)
diversity_loss = diversity_loss/(torch.sum(index_mask))/(torch.sum(index_mask)-1)



  
  
    
    
  
  


p.s. The torch.index_select in below is just for training set preparation so I do not change.

  
      github.com
  
  
    pcshih/pytorch-VSLUD/blob/b3ed7aba9332d2e0a21a66e84eae3654c9e254af/training_set_preparation.py#L42-L47

gt_summary = torch.from_numpy(dataset[key]["gtsummary"][...]).to(device)
column_index = gt_summary.nonzero().view(-1)
feature_summary_cuda = torch.from_numpy(dataset[key]["features"][...]).to(device)
feature_summary_cuda = feature_summary_cuda.transpose(1,0).view(1,1024,1,feature_summary_cuda.shape[0])
feature_summary_cuda = torch.index_select(feature_summary_cuda, 3, column_index)
attributes["summary_features"] = feature_summary_cuda; #print(torch.isnan(feature_summary_cuda).nonzero().view(-1))



  
  
    
    
  
  


Thank you very much, Tony-Y.