| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| """ |
| Base Encoder to create multi-level conv features and positional embedding. |
| """ |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from misc.detr_utils.misc import NestedTensor |
| from .position_encoding import PositionEmbeddingSine |
|
|
|
|
| class BaseEncoder(nn.Module): |
| def __init__(self, num_feature_levels, vf_dim, hidden_dim): |
| super(BaseEncoder, self).__init__() |
| self.pos_embed = PositionEmbeddingSine(hidden_dim//2, normalize=True) |
| self.num_feature_levels = num_feature_levels |
| self.hidden_dim = hidden_dim |
|
|
| if num_feature_levels > 1: |
| input_proj_list = [] |
| in_channels = vf_dim |
| input_proj_list.append(nn.Sequential( |
| nn.Conv1d(in_channels, hidden_dim, kernel_size=1), |
| nn.GroupNorm(32, hidden_dim), |
| )) |
| for _ in range(num_feature_levels - 1): |
| input_proj_list.append(nn.Sequential( |
| nn.Conv1d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1), |
| nn.GroupNorm(32, hidden_dim), |
| )) |
| in_channels = hidden_dim |
| self.input_proj = nn.ModuleList(input_proj_list) |
| else: |
| self.input_proj = nn.ModuleList([ |
| nn.Sequential( |
| nn.Conv2d(vf_dim, hidden_dim, kernel_size=1), |
| nn.GroupNorm(32, hidden_dim), |
| )]) |
|
|
| for proj in self.input_proj: |
| nn.init.xavier_uniform_(proj[0].weight, gain=1) |
| nn.init.constant_(proj[0].bias, 0) |
|
|
| def forward(self, vf, mask, duration): |
| |
| vf = vf.transpose(1, 2) |
| vf_nt = NestedTensor(vf, mask, duration) |
| pos0 = self.pos_embed(vf_nt) |
|
|
| srcs = [] |
| masks = [] |
| poses = [] |
|
|
| src0, mask0 = vf_nt.decompose() |
| srcs.append(self.input_proj[0](src0)) |
| masks.append(mask0) |
| poses.append(pos0) |
| assert mask is not None |
|
|
| for l in range(1, self.num_feature_levels): |
| if l == 1: |
| src = self.input_proj[l](vf_nt.tensors) |
| else: |
| src = self.input_proj[l](srcs[-1]) |
| m = vf_nt.mask |
| mask = F.interpolate(m[None].float(), size=src.shape[-1:]).to(torch.bool)[0] |
| pos_l = self.pos_embed(NestedTensor(src, mask, duration)).to(src.dtype) |
| srcs.append(src) |
| masks.append(mask) |
| poses.append(pos_l) |
| return srcs, masks, poses |
|
|
| def build_base_encoder(args): |
| base_encoder = BaseEncoder(args.num_feature_levels, args.feature_dim, args.hidden_dim) |
| return base_encoder |
|
|
