289 lines
9.7 KiB
Python
289 lines
9.7 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class Q2VRankerStage1(nn.Module):
|
|
"""
|
|
Used to calculate the qv_ctx_score with query embedding and multi anchor context embeddings as input.
|
|
The qv_ctx_score is used to pre-rank and retain top-k related anchors.
|
|
"""
|
|
|
|
def __init__(self, nscales, hidden_dim):
|
|
super().__init__()
|
|
self.fc = nn.Linear(hidden_dim, hidden_dim)
|
|
self.nscales = nscales
|
|
|
|
def forward(self, ctx_feats, qfeat):
|
|
qfeat = self.fc(qfeat)
|
|
qv_ctx_scores = list()
|
|
for i in range(self.nscales):
|
|
score = torch.einsum('bld,bd->bl',
|
|
F.normalize(ctx_feats[i], p=2, dim=2),
|
|
F.normalize(qfeat, p=2, dim=1))
|
|
qv_ctx_scores.append(score)
|
|
|
|
return qv_ctx_scores
|
|
|
|
|
|
class V2QRankerStage1(nn.Module):
|
|
"""
|
|
Used to calculate the vq_ctx_score with anchor context embeddings and multi query embeddings as input.
|
|
"""
|
|
|
|
def __init__(self, nscales, hidden_dim):
|
|
super().__init__()
|
|
self.fc = nn.Linear(hidden_dim, hidden_dim)
|
|
self.nscales = nscales
|
|
|
|
def forward(self, ctx_feats, qfeat):
|
|
vq_ctx_scores = list()
|
|
for i in range(self.nscales):
|
|
score = torch.einsum(
|
|
'bld,bd->bl', F.normalize(self.fc(ctx_feats[i]), p=2, dim=2),
|
|
F.normalize(qfeat, p=2, dim=1))
|
|
vq_ctx_scores.append(score)
|
|
|
|
return vq_ctx_scores
|
|
|
|
|
|
class Q2VRankerStage2(nn.Module):
|
|
"""
|
|
Used to calculate the qv_ctn_score with query embedding and video sequence embedding as input.
|
|
The qv_ctn_score is used to re-rank anchors.
|
|
"""
|
|
|
|
def __init__(self, nscales, hidden_dim, snippet_length=10):
|
|
super().__init__()
|
|
self.nscales = nscales
|
|
self.snippet_length = snippet_length
|
|
self.qfc = nn.Linear(hidden_dim, hidden_dim)
|
|
self.encoder = V2VAttention()
|
|
|
|
def forward(self, vfeats, qfeat, hit_indices, qv_ctx_scores):
|
|
qfeat = self.qfc(qfeat)
|
|
|
|
qv_ctn_scores = list()
|
|
qv_merge_scores = list()
|
|
|
|
_, L, D = vfeats.size()
|
|
ctn_feats = list()
|
|
for i in range(self.nscales):
|
|
anchor_length = self.snippet_length * 2**i
|
|
assert L // anchor_length == qv_ctx_scores[i].size(1)
|
|
qv_ctx_score = torch.index_select(qv_ctx_scores[i], 1,
|
|
hit_indices[i])
|
|
|
|
ctn_feat = vfeats.view(L // anchor_length, anchor_length,
|
|
D).detach()
|
|
ctn_feat = torch.index_select(ctn_feat, 0, hit_indices[i])
|
|
ctn_feat = self.encoder(
|
|
ctn_feat,
|
|
torch.ones(ctn_feat.size()[:2], device=ctn_feat.device))
|
|
ctn_feats.append(ctn_feat)
|
|
|
|
qv_ctn_score = torch.einsum(
|
|
'bkld,bd->bkl', F.normalize(ctn_feat.unsqueeze(0), p=2, dim=3),
|
|
F.normalize(qfeat, p=2, dim=1))
|
|
qv_ctn_score, _ = torch.max(qv_ctn_score, dim=2)
|
|
qv_ctn_scores.append(qv_ctn_score)
|
|
qv_merge_scores.append(qv_ctx_score + qv_ctn_score)
|
|
|
|
return qv_merge_scores, qv_ctn_scores, ctn_feats
|
|
|
|
|
|
class V2QRankerStage2(nn.Module):
|
|
"""
|
|
Used to calculate the vq_ctn_score with anchor content embeddings and multi query embeddings as input.
|
|
"""
|
|
|
|
def __init__(self, nscales, hidden_dim):
|
|
super().__init__()
|
|
self.fc = nn.Linear(hidden_dim, hidden_dim)
|
|
self.nscales = nscales
|
|
|
|
def forward(self, ctn_feats, qfeat):
|
|
vq_ctn_scores = list()
|
|
for i in range(self.nscales):
|
|
score = torch.einsum(
|
|
'bkld,bd->bkl',
|
|
F.normalize(self.fc(ctn_feats[i]).unsqueeze(0), p=2, dim=3),
|
|
F.normalize(qfeat, p=2, dim=1))
|
|
score = torch.mean(score, dim=2)
|
|
vq_ctn_scores.append(score)
|
|
|
|
return vq_ctn_scores
|
|
|
|
|
|
class V2VAttention(nn.Module):
|
|
"""
|
|
Self-attention encoder for anchor frame sequence to encode intra-anchor knowledge.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.posemb = PositionEncoding(max_len=400, dim=512, dropout=0.0)
|
|
self.encoder = MultiHeadAttention(dim=512, n_heads=8, dropout=0.1)
|
|
self.dropout = nn.Dropout(0.0)
|
|
|
|
def forward(self, video_feats, video_masks):
|
|
mask = torch.einsum('bm,bn->bmn', video_masks,
|
|
video_masks).unsqueeze(1)
|
|
residual = video_feats
|
|
video_feats = video_feats + self.posemb(video_feats)
|
|
out = self.encoder(
|
|
query=video_feats, key=video_feats, value=video_feats, mask=mask)
|
|
video_feats = self.dropout(residual
|
|
+ out) * video_masks.unsqueeze(2).float()
|
|
return video_feats
|
|
|
|
|
|
class BboxRegressor(nn.Module):
|
|
"""
|
|
Predict the offset of bounding box for each candidate anchor.
|
|
"""
|
|
|
|
def __init__(self, hidden_dim, enable_stage2=False):
|
|
super().__init__()
|
|
self.fc_ctx = nn.Linear(hidden_dim, hidden_dim)
|
|
self.fc_q = nn.Linear(hidden_dim, hidden_dim)
|
|
|
|
if enable_stage2:
|
|
self.fc_ctn = nn.Linear(hidden_dim, hidden_dim)
|
|
self.attn = SelfAttention(hidden_dim)
|
|
self.predictor = nn.Sequential(
|
|
nn.Linear(2 * hidden_dim, hidden_dim), nn.ReLU(),
|
|
nn.Linear(hidden_dim, 2))
|
|
else:
|
|
self.predictor = nn.Sequential(
|
|
nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
|
|
nn.Linear(hidden_dim, 2))
|
|
self.enable_stage2 = enable_stage2
|
|
|
|
def forward(self, ctx_feats, ctn_feats, qfeat):
|
|
qfeat = self.fc_q(qfeat)
|
|
|
|
ctx_feats = torch.cat(ctx_feats, dim=1)
|
|
ctx_fuse_feats = F.relu(self.fc_ctx(ctx_feats)) * F.relu(
|
|
qfeat.unsqueeze(1))
|
|
|
|
if self.enable_stage2 and ctn_feats:
|
|
ctn_fuse_feats = list()
|
|
for i in range(len(ctn_feats)):
|
|
out = F.relu(self.fc_ctn(ctn_feats[i]).unsqueeze(0)) * F.relu(
|
|
qfeat.unsqueeze(1).unsqueeze(1))
|
|
out = self.attn(out)
|
|
ctn_fuse_feats.append(out)
|
|
ctn_fuse_feats = torch.cat(ctn_fuse_feats, dim=1)
|
|
fuse_feats = torch.cat([ctx_fuse_feats, ctn_fuse_feats], dim=-1)
|
|
else:
|
|
fuse_feats = ctx_fuse_feats
|
|
|
|
out = self.predictor(fuse_feats)
|
|
return out
|
|
|
|
|
|
class SelfAttention(nn.Module):
|
|
"""
|
|
Obtain pooled features by self-attentive pooling.
|
|
"""
|
|
|
|
def __init__(self, hidden_dim):
|
|
super().__init__()
|
|
self.fc1 = nn.Linear(hidden_dim, hidden_dim // 2)
|
|
self.relu = nn.ReLU()
|
|
self.fc2 = nn.Linear(hidden_dim // 2, 1)
|
|
|
|
def forward(self, x):
|
|
att = self.fc2(self.relu(self.fc1(x))).squeeze(3)
|
|
att = F.softmax(att, dim=2).unsqueeze(3)
|
|
out = torch.sum(x * att, dim=2)
|
|
return out
|
|
|
|
|
|
class PositionEncoding(nn.Module):
|
|
"""
|
|
An implementation of trainable positional embedding which is added to
|
|
sequence features to inject time/position information.
|
|
|
|
Args:
|
|
max_len: The max number of trainable positional embeddings.
|
|
dim: the dimension of positional embedding.
|
|
"""
|
|
|
|
def __init__(self, max_len, dim, dropout=0.0):
|
|
super(PositionEncoding, self).__init__()
|
|
|
|
self.embed = nn.Embedding(max_len, dim)
|
|
self.relu = nn.ReLU()
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
def forward(self, x):
|
|
batch_size, seq_len = x.shape[:2]
|
|
pos_ids = torch.arange(seq_len, dtype=torch.long, device=x.device)
|
|
pos_ids = pos_ids.unsqueeze(0).repeat(batch_size, 1)
|
|
pos_emb = self.dropout(self.relu(self.embed(pos_ids)))
|
|
|
|
return pos_emb
|
|
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
"""
|
|
An implementation of multi-head attention module, as described in
|
|
'Attention Is All You Need <https://arxiv.org/abs/1706.03762>'
|
|
|
|
Args:
|
|
dim: the dimension of features of hidden layers.
|
|
n_heads: the number of head.
|
|
"""
|
|
|
|
def __init__(self, dim, n_heads, dropout=0.0):
|
|
super(MultiHeadAttention, self).__init__()
|
|
|
|
self.dim = dim
|
|
self.n_heads = n_heads
|
|
self.head_dim = dim // n_heads
|
|
|
|
self.to_q = nn.Linear(dim, dim)
|
|
self.to_k = nn.Linear(dim, dim)
|
|
self.to_v = nn.Linear(dim, dim)
|
|
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.softmax = nn.Softmax(dim=-1)
|
|
|
|
def transpose_for_scores(self, x):
|
|
new_x_shape = x.size()[:-1] + (self.n_heads, self.head_dim)
|
|
x = x.view(*new_x_shape)
|
|
return x.permute(0, 2, 1, 3) # (N, nh, L, dh)
|
|
|
|
def forward(self, query, key, value, mask):
|
|
q = self.to_q(query)
|
|
k = self.to_k(key)
|
|
v = self.to_v(value)
|
|
|
|
q_trans = self.transpose_for_scores(q)
|
|
k_trans = self.transpose_for_scores(k)
|
|
v_trans = self.transpose_for_scores(v)
|
|
|
|
att = torch.matmul(q_trans, k_trans.transpose(-1,
|
|
-2)) # (N, nh, Lq, L)
|
|
att = att / math.sqrt(self.head_dim)
|
|
att = mask_logits(att, mask)
|
|
att = self.softmax(att)
|
|
att = self.dropout(att)
|
|
|
|
ctx_v = torch.matmul(att, v_trans) # (N, nh, Lq, dh)
|
|
ctx_v = ctx_v.permute(0, 2, 1, 3).contiguous() # (N, Lq, nh, dh)
|
|
shape = ctx_v.size()[:-2] + (self.dim, )
|
|
ctx_v = ctx_v.view(*shape) # (N, Lq, D)
|
|
return ctx_v
|
|
|
|
|
|
def mask_logits(inputs, mask, mask_value=-1e30):
|
|
mask = mask.type(torch.float32)
|
|
return inputs + (1.0 - mask) * mask_value
|