351 lines
14 KiB
Python
Executable File
351 lines
14 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from configs.feature_sizes import FEATURE_SIZES
|
|
from model.model_builder import META_ARCHITECTURES as registry
|
|
|
|
from .attn import AttentionLayer, FullAttention, ProbAttention
|
|
from .decoder import Decoder, DecoderLayer
|
|
from .PositionalEncoding import (FixedPositionalEncoding,
|
|
LearnedPositionalEncoding)
|
|
from .Transformer import TransformerModel
|
|
|
|
|
|
# @registry.register('TRANSFORMER')
|
|
@registry.register('OADTR')
|
|
class OADTR(nn.Module):
|
|
def __init__(self, cfg, use_representation=True,
|
|
conv_patch_representation=False,
|
|
positional_encoding_type="learned"
|
|
):
|
|
super(OADTR, self).__init__()
|
|
|
|
self.img_dim = cfg["window_size"] # img_dim
|
|
self.out_dim = cfg["num_classes"]
|
|
|
|
self.patch_dim = cfg["patch_dim"]
|
|
self.query_num = cfg["query_num"]
|
|
|
|
self.embedding_dim = cfg["embedding_dim"]
|
|
self.num_heads = cfg["num_heads"]
|
|
self.num_layers = cfg["num_layers"]
|
|
self.hidden_dim = cfg["hidden_dim"]
|
|
|
|
self.dropout_rate = cfg["dropout"]
|
|
self.attn_dropout_rate = cfg["attn_dropout_rate"]
|
|
|
|
self.conv_patch_representation = conv_patch_representation
|
|
|
|
self.decoder_embedding_dim = cfg["decoder_embedding_dim"]
|
|
self.decoder_num_heads = cfg["decoder_num_heads"]
|
|
self.decoder_layers = cfg["decoder_layers"]
|
|
self.decoder_embedding_dim_out = cfg["decoder_embedding_dim_out"]
|
|
self.decoder_attn_dropout_rate = cfg["decoder_attn_dropout_rate"]
|
|
|
|
self.use_flow = not cfg['no_flow']
|
|
self.use_rgb = not cfg['no_rgb']
|
|
|
|
self.num_channels = 0
|
|
if self.use_rgb:
|
|
self.num_channels += FEATURE_SIZES[cfg['rgb_type']]
|
|
if self.use_flow:
|
|
self.num_channels += FEATURE_SIZES[cfg['flow_type']]
|
|
|
|
assert self.embedding_dim % self.num_heads == 0
|
|
assert self.img_dim % self.patch_dim == 0
|
|
|
|
self.num_patches = int(self.img_dim // self.patch_dim)
|
|
self.seq_length = self.num_patches + 1
|
|
# self.seq_length = self.num_patches + self.img_dim
|
|
self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels
|
|
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim))
|
|
# self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim))
|
|
|
|
self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
|
|
if positional_encoding_type == "learned":
|
|
self.position_encoding = LearnedPositionalEncoding(
|
|
self.seq_length, self.embedding_dim, self.seq_length
|
|
)
|
|
elif positional_encoding_type == "fixed":
|
|
self.position_encoding = FixedPositionalEncoding(
|
|
self.embedding_dim,
|
|
)
|
|
print('position encoding :', positional_encoding_type)
|
|
|
|
self.pe_dropout = nn.Dropout(p=self.dropout_rate)
|
|
|
|
self.encoder = TransformerModel(
|
|
self.embedding_dim,
|
|
self.num_layers,
|
|
self.num_heads,
|
|
self.hidden_dim,
|
|
self.dropout_rate,
|
|
self.attn_dropout_rate,
|
|
)
|
|
self.pre_head_ln = nn.LayerNorm(self.embedding_dim)
|
|
|
|
d_model = self.decoder_embedding_dim
|
|
use_representation = False # False
|
|
if use_representation:
|
|
self.mlp_head = nn.Sequential(
|
|
nn.Linear(self.embedding_dim , self.hidden_dim//2),
|
|
# nn.Tanh(),
|
|
nn.ReLU(),
|
|
nn.Linear(self.hidden_dim//2, self.out_dim),
|
|
)
|
|
else:
|
|
self.mlp_head = nn.Linear(self.embedding_dim + d_model, self.out_dim)
|
|
|
|
if self.conv_patch_representation:
|
|
# self.conv_x = nn.Conv2d(
|
|
# self.num_channels,
|
|
# self.embedding_dim,
|
|
# kernel_size=(self.patch_dim, self.patch_dim),
|
|
# stride=(self.patch_dim, self.patch_dim),
|
|
# padding=self._get_padding(
|
|
# 'VALID', (self.patch_dim, self.patch_dim),
|
|
# ),
|
|
# )
|
|
self.conv_x = nn.Conv1d(
|
|
self.num_channels,
|
|
self.embedding_dim,
|
|
kernel_size=self.patch_dim,
|
|
stride=self.patch_dim,
|
|
padding=self._get_padding(
|
|
'VALID', (self.patch_dim),
|
|
),
|
|
)
|
|
else:
|
|
self.conv_x = None
|
|
|
|
self.to_cls_token = nn.Identity()
|
|
|
|
# Decoder
|
|
factor = 1 # 5
|
|
dropout = self.decoder_attn_dropout_rate
|
|
# d_model = self.decoder_embedding_dim
|
|
n_heads = self.decoder_num_heads
|
|
d_layers = self.decoder_layers
|
|
d_ff = self.decoder_embedding_dim_out # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None
|
|
activation = 'gelu' # 'gelu'
|
|
self.decoder = Decoder(
|
|
[
|
|
DecoderLayer(
|
|
AttentionLayer(FullAttention(True, factor, attention_dropout=dropout), # True
|
|
d_model, n_heads), # ProbAttention FullAttention
|
|
AttentionLayer(FullAttention(False, factor, attention_dropout=dropout), # False
|
|
d_model, n_heads),
|
|
d_model,
|
|
d_ff,
|
|
dropout=dropout,
|
|
activation=activation,
|
|
)
|
|
for l in range(d_layers)
|
|
],
|
|
norm_layer=torch.nn.LayerNorm(d_model)
|
|
)
|
|
self.decoder_cls_token = nn.Parameter(torch.zeros(1, self.query_num, d_model))
|
|
if positional_encoding_type == "learned":
|
|
self.decoder_position_encoding = LearnedPositionalEncoding(
|
|
self.query_num, self.embedding_dim, self.query_num
|
|
)
|
|
elif positional_encoding_type == "fixed":
|
|
self.decoder_position_encoding = FixedPositionalEncoding(
|
|
self.embedding_dim,
|
|
)
|
|
print('position decoding :', positional_encoding_type)
|
|
self.classifier = nn.Linear(d_model, self.out_dim)
|
|
self.after_dropout = nn.Dropout(p=self.dropout_rate)
|
|
# self.merge_fc = nn.Linear(d_model, 1)
|
|
# self.merge_sigmoid = nn.Sigmoid()
|
|
|
|
|
|
def forward(self, sequence_input_rgb, sequence_input_flow):
|
|
if self.use_rgb and self.use_flow:
|
|
x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
|
|
elif self.use_rgb:
|
|
x = sequence_input_rgb
|
|
elif self.use_flow:
|
|
x = sequence_input_flow
|
|
|
|
x = self.linear_encoding(x) # B, 128, 1204
|
|
cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024
|
|
# x = torch.cat((cls_tokens, x), dim=1)
|
|
x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024
|
|
try:
|
|
x = self.position_encoding(x) # B, seq+1, 1024
|
|
except:
|
|
import pdb; pdb.set_trace()
|
|
x = self.pe_dropout(x) # not delete
|
|
|
|
# apply transformer
|
|
x = self.encoder(x)
|
|
x = self.pre_head_ln(x) # B, seq+1, 1024
|
|
|
|
# decoder
|
|
decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1) # [B, 8, 1024]
|
|
# decoder_cls_token = self.after_dropout(decoder_cls_token) # add
|
|
# decoder_cls_token = self.decoder_position_encoding(decoder_cls_token) # [128, 8, 1024]
|
|
dec = self.decoder(decoder_cls_token, x) # [128, 8, 1024]
|
|
dec = self.after_dropout(dec) # add
|
|
|
|
# merge_atte = self.merge_sigmoid(self.merge_fc(dec)) # [128, 8, 1]
|
|
# dec_for_token = (merge_atte*dec).sum(dim=1) # [128, 1024]
|
|
# dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001)
|
|
# *** RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 1024x2)
|
|
dec_for_token = dec.mean(dim=1) # [64, 1024]
|
|
|
|
# dec_for_token = dec.max(dim=1)[0]
|
|
dec_cls_out = self.classifier(dec) # [B, 8, NUM_CLASS?]
|
|
# set_trace()
|
|
# x = self.to_cls_token(x[:, 0])
|
|
|
|
x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1) # [B, 2048]
|
|
x = self.mlp_head(x) # [B, NUM_CLASSES]
|
|
|
|
# original OADTR
|
|
# return x, dec_cls_out
|
|
|
|
# x = F.log_softmax(x, dim=-1)
|
|
out_dict = {}
|
|
out_dict['logits'] = x.unsqueeze(1) # [B, 1, NUM_CLASSES] ex) [64, 1, 2]
|
|
out_dict['dec_logits'] = dec_cls_out # [B, 8, NUM_CLASSES] ex) [64, 8, 2]
|
|
|
|
# out_dict['logits'] = x #x.unsqueeze(1)
|
|
return out_dict
|
|
|
|
def _get_padding(self, padding_type, kernel_size):
|
|
assert padding_type in ['SAME', 'VALID']
|
|
if padding_type == 'SAME':
|
|
_list = [(k - 1) // 2 for k in kernel_size]
|
|
return tuple(_list)
|
|
return tuple(0 for _ in kernel_size)
|
|
|
|
@registry.register('Transformer')
|
|
class ViTEnc(nn.Module):
|
|
def __init__(self, cfg, use_representation=True,
|
|
conv_patch_representation=False,
|
|
positional_encoding_type="learned"
|
|
):
|
|
super(ViTEnc, self).__init__()
|
|
self.img_dim = cfg["window_size"]
|
|
self.out_dim = cfg["num_classes"]
|
|
self.embedding_dim = cfg["embedding_dim"]
|
|
self.patch_dim = cfg["patch_dim"]
|
|
self.num_heads = cfg["num_heads"]
|
|
self.num_layers = cfg["num_layers"]
|
|
self.hidden_dim = cfg["hidden_dim"]
|
|
self.dropout_rate = cfg["dropout"]
|
|
self.use_flow = not cfg['no_flow']
|
|
self.use_rgb = not cfg['no_rgb']
|
|
self.num_channels= 0
|
|
if self.use_rgb:
|
|
self.num_channels += FEATURE_SIZES[cfg['rgb_type']]
|
|
if self.use_flow:
|
|
self.num_channels += FEATURE_SIZES[cfg['flow_type']]
|
|
self.attn_dropout_rate = cfg["attn_dropout_rate"]
|
|
assert self.embedding_dim % self.num_heads == 0
|
|
assert self.img_dim % self.patch_dim == 0
|
|
self.conv_patch_representation = conv_patch_representation
|
|
self.num_patches = int(self.img_dim // self.patch_dim)
|
|
self.seq_length = self.num_patches + 1
|
|
# self.seq_length = self.num_patches + self.img_dim
|
|
self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels
|
|
self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim))
|
|
# self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim))
|
|
|
|
self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
|
|
if positional_encoding_type == "learned":
|
|
self.position_encoding = LearnedPositionalEncoding(
|
|
self.seq_length, self.embedding_dim, self.seq_length
|
|
)
|
|
elif positional_encoding_type == "fixed":
|
|
self.position_encoding = FixedPositionalEncoding(
|
|
self.embedding_dim,
|
|
)
|
|
print('position encoding :', positional_encoding_type)
|
|
|
|
self.pe_dropout = nn.Dropout(p=self.dropout_rate)
|
|
|
|
self.encoder = TransformerModel(
|
|
self.embedding_dim,
|
|
self.num_layers,
|
|
self.num_heads,
|
|
self.hidden_dim,
|
|
self.dropout_rate,
|
|
self.attn_dropout_rate,
|
|
)
|
|
self.pre_head_ln = nn.LayerNorm(self.embedding_dim)
|
|
|
|
use_representation = False # False
|
|
if use_representation:
|
|
self.mlp_head = nn.Sequential(
|
|
nn.Linear(self.embedding_dim , self.hidden_dim//2),
|
|
# nn.Tanh(),
|
|
nn.ReLU(),
|
|
nn.Linear(self.hidden_dim//2, self.out_dim),
|
|
)
|
|
else:
|
|
self.mlp_head = nn.Linear(self.embedding_dim, self.out_dim)
|
|
|
|
if self.conv_patch_representation:
|
|
# self.conv_x = nn.Conv2d(
|
|
# self.num_channels,
|
|
# self.embedding_dim,
|
|
# kernel_size=(self.patch_dim, self.patch_dim),
|
|
# stride=(self.patch_dim, self.patch_dim),
|
|
# padding=self._get_padding(
|
|
# 'VALID', (self.patch_dim, self.patch_dim),
|
|
# ),
|
|
# )
|
|
self.conv_x = nn.Conv1d(
|
|
self.num_channels,
|
|
self.embedding_dim,
|
|
kernel_size=self.patch_dim,
|
|
stride=self.patch_dim,
|
|
padding=self._get_padding(
|
|
'VALID', (self.patch_dim),
|
|
),
|
|
)
|
|
else:
|
|
self.conv_x = None
|
|
|
|
self.to_cls_token = nn.Identity()
|
|
|
|
|
|
def forward(self, sequence_input_rgb, sequence_input_flow):
|
|
if self.use_rgb and self.use_flow:
|
|
x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
|
|
elif self.use_rgb:
|
|
x = sequence_input_rgb
|
|
elif self.use_flow:
|
|
x = sequence_input_flow
|
|
|
|
x = self.linear_encoding(x)
|
|
cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024
|
|
# x = torch.cat((cls_tokens, x), dim=1)
|
|
x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024
|
|
|
|
x = self.position_encoding(x) # B, seq+1, 1024
|
|
x = self.pe_dropout(x) # not delete
|
|
|
|
# apply transformer
|
|
x = self.encoder(x)
|
|
x = self.pre_head_ln(x) # B, seq+1, 1024
|
|
|
|
x = self.to_cls_token(x[:, 0]) # B, 1024
|
|
# x = self.to_cls_token(x[:,0:self.img_dim]) # B, 1024
|
|
x = self.mlp_head(x)
|
|
# x = F.log_softmax(x, dim=-1)
|
|
out_dict = {}
|
|
out_dict['logits'] = x.unsqueeze(1)
|
|
# out_dict['logits'] = x #x.unsqueeze(1)
|
|
return out_dict
|
|
|
|
def _get_padding(self, padding_type, kernel_size):
|
|
assert padding_type in ['SAME', 'VALID']
|
|
if padding_type == 'SAME':
|
|
_list = [(k - 1) // 2 for k in kernel_size]
|
|
return tuple(_list)
|
|
return tuple(0 for _ in kernel_size)
|