forked from cmhungsteve/TA3N
-
Notifications
You must be signed in to change notification settings - Fork 0
/
TRNmodule.py
86 lines (74 loc) · 4.03 KB
/
TRNmodule.py
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
import torch
import torch.nn as nn
import numpy as np
from math import ceil
class RelationModule(torch.nn.Module):
# this is the naive implementation of the n-frame relation module, as num_frames == num_frames_relation
def __init__(self, img_feature_dim, num_bottleneck, num_frames):
super(RelationModule, self).__init__()
self.num_frames = num_frames
self.img_feature_dim = img_feature_dim
self.num_bottleneck = num_bottleneck
self.classifier = self.fc_fusion()
def fc_fusion(self):
# naive concatenate
classifier = nn.Sequential(
nn.ReLU(),
nn.Linear(self.num_frames * self.img_feature_dim, self.num_bottleneck),
nn.ReLU(),
)
return classifier
def forward(self, input):
input = input.view(input.size(0), self.num_frames*self.img_feature_dim)
input = self.classifier(input)
return input
class RelationModuleMultiScale(torch.nn.Module):
# Temporal Relation module in multiply scale, suming over [2-frame relation, 3-frame relation, ..., n-frame relation]
def __init__(self, img_feature_dim, num_bottleneck, num_frames):
super(RelationModuleMultiScale, self).__init__()
self.subsample_num = 3 # how many relations selected to sum up
self.img_feature_dim = img_feature_dim
self.scales = [i for i in range(num_frames, 1, -1)] # generate the multiple frame relations
self.relations_scales = []
self.subsample_scales = []
for scale in self.scales:
relations_scale = self.return_relationset(num_frames, scale)
self.relations_scales.append(relations_scale)
self.subsample_scales.append(min(self.subsample_num, len(relations_scale))) # how many samples of relation to select in each forward pass
# self.num_class = num_class
self.num_frames = num_frames
self.fc_fusion_scales = nn.ModuleList() # high-tech modulelist
for i in range(len(self.scales)):
scale = self.scales[i]
fc_fusion = nn.Sequential(
nn.ReLU(),
nn.Linear(scale * self.img_feature_dim, num_bottleneck),
nn.ReLU(),
)
self.fc_fusion_scales += [fc_fusion]
print('Multi-Scale Temporal Relation Network Module in use', ['%d-frame relation' % i for i in self.scales])
def forward(self, input):
# the first one is the largest scale
act_scale_1 = input[:, self.relations_scales[0][0] , :]
act_scale_1 = act_scale_1.view(act_scale_1.size(0), self.scales[0] * self.img_feature_dim)
act_scale_1 = self.fc_fusion_scales[0](act_scale_1)
act_scale_1 = act_scale_1.unsqueeze(1) # add one dimension for the later concatenation
act_all = act_scale_1.clone()
for scaleID in range(1, len(self.scales)):
act_relation_all = torch.zeros_like(act_scale_1)
# iterate over the scales
num_total_relations = len(self.relations_scales[scaleID])
num_select_relations = self.subsample_scales[scaleID]
idx_relations_evensample = [int(ceil(i * num_total_relations / num_select_relations)) for i in range(num_select_relations)]
#for idx in idx_relations_randomsample:
for idx in idx_relations_evensample:
act_relation = input[:, self.relations_scales[scaleID][idx], :]
act_relation = act_relation.view(act_relation.size(0), self.scales[scaleID] * self.img_feature_dim)
act_relation = self.fc_fusion_scales[scaleID](act_relation)
act_relation = act_relation.unsqueeze(1) # add one dimension for the later concatenation
act_relation_all += act_relation
act_all = torch.cat((act_all, act_relation_all), 1)
return act_all
def return_relationset(self, num_frames, num_frames_relation):
import itertools
return list(itertools.combinations([i for i in range(num_frames)], num_frames_relation))