1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
|
import copy
import torch
import torch.nn as nn
from models.layers import BasicConv1d, FocalConv2dBlock
class FrameLevelPartFeatureExtractor(nn.Module):
def __init__(
self,
in_channels: int = 3,
feature_channels: int = 32,
kernel_sizes: tuple[tuple, ...] = ((5, 3), (3, 3), (3, 3)),
paddings: tuple[tuple, ...] = ((2, 1), (1, 1), (1, 1)),
halving: tuple[int, ...] = (0, 2, 3)
):
super().__init__()
num_blocks = len(kernel_sizes)
out_channels = [feature_channels * 2 ** i for i in range(num_blocks)]
in_channels = [in_channels] + out_channels[:-1]
use_pools = [True] * (num_blocks - 1) + [False]
params = (in_channels, out_channels, kernel_sizes,
paddings, halving, use_pools)
self.fconv_blocks = nn.ModuleList([
FocalConv2dBlock(*_params) for _params in zip(*params)
])
def forward(self, x):
# Flatten frames in all batches
n, t, c, h, w = x.size()
x = x.view(n * t, c, h, w)
for fconv_block in self.fconv_blocks:
x = fconv_block(x)
return x
class TemporalFeatureAggregator(nn.Module):
def __init__(
self,
in_channels: int,
squeeze_ratio: int = 4,
num_part: int = 16
):
super().__init__()
hidden_dim = in_channels // squeeze_ratio
self.num_part = num_part
# MTB1
conv3x1 = nn.Sequential(
BasicConv1d(in_channels, hidden_dim, kernel_size=3, padding=1),
nn.LeakyReLU(inplace=True),
BasicConv1d(hidden_dim, in_channels, kernel_size=1, padding=0)
)
self.conv1d3x1 = self._parted(conv3x1)
self.avg_pool3x1 = nn.AvgPool1d(kernel_size=3, stride=1, padding=1)
self.max_pool3x1 = nn.MaxPool1d(kernel_size=3, stride=1, padding=1)
# MTB2
conv3x3 = nn.Sequential(
BasicConv1d(in_channels, hidden_dim, kernel_size=3, padding=1),
nn.LeakyReLU(inplace=True),
BasicConv1d(hidden_dim, in_channels, kernel_size=3, padding=1)
)
self.conv1d3x3 = self._parted(conv3x3)
self.avg_pool3x3 = nn.AvgPool1d(kernel_size=5, stride=1, padding=2)
self.max_pool3x3 = nn.MaxPool1d(kernel_size=5, stride=1, padding=2)
def _parted(self, module: nn.Module):
"""Duplicate module `part_num` times."""
return nn.ModuleList([copy.deepcopy(module)
for _ in range(self.num_part)])
def forward(self, x):
# p, n, t, c
x = x.transpose(2, 3)
p, n, c, t = x.size()
feature = x.split(1, dim=0)
feature = [f.squeeze(0) for f in feature]
x = x.view(-1, c, t)
# MTB1: ConvNet1d & Sigmoid
logits3x1 = torch.stack(
[conv(f) for conv, f in zip(self.conv1d3x1, feature)]
)
scores3x1 = torch.sigmoid(logits3x1)
# MTB1: Template Function
feature3x1 = self.avg_pool3x1(x) + self.max_pool3x1(x)
feature3x1 = feature3x1.view(p, n, c, t)
feature3x1 = feature3x1 * scores3x1
# MTB2: ConvNet1d & Sigmoid
logits3x3 = torch.stack(
[conv(f) for conv, f in zip(self.conv1d3x3, feature)]
)
scores3x3 = torch.sigmoid(logits3x3)
# MTB2: Template Function
feature3x3 = self.avg_pool3x3(x) + self.max_pool3x3(x)
feature3x3 = feature3x3.view(p, n, c, t)
feature3x3 = feature3x3 * scores3x3
# Temporal Pooling
ret = (feature3x1 + feature3x3).max(-1)[0]
return ret
class PartNet(nn.Module):
def __init__(
self,
in_channels: int = 128,
embedding_dims: int = 256,
num_parts: int = 16,
squeeze_ratio: int = 4,
):
super().__init__()
self.num_part = num_parts
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.max_pool = nn.AdaptiveMaxPool2d(1)
self.tfa = TemporalFeatureAggregator(
in_channels, squeeze_ratio, self.num_part
)
self.fc_mat = nn.Parameter(
torch.empty(num_parts, in_channels, embedding_dims)
)
def forward(self, x):
n, t, c, h, w = x.size()
x = x.view(n * t, c, h, w)
# n * t x c x h x w
# Horizontal Pooling
_, c, h, w = x.size()
split_size = h // self.num_part
x = x.split(split_size, dim=2)
x = [self.avg_pool(x_) + self.max_pool(x_) for x_ in x]
x = [x_.view(n, t, c) for x_ in x]
x = torch.stack(x)
# p, n, t, c
x = self.tfa(x)
# p, n, c
x = x @ self.fc_mat
# p, n, d
return x
|
