SMIE/model.py at main · YujieOuO/SMIE · GitHub

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
from config import *

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import repeat,rearrange
import random

def init_weights(m):
    class_name=m.__class__.__name__

    if "Linear" in class_name:
        for name, param in m.named_parameters():
            if 'weight' in name:
                nn.init.xavier_normal_(param.data)
            elif 'bias' in name:
                nn.init.constant_(param.data, 0.0)

class GlobalDiscriminator(nn.Module):

    @ex.capture
    def __init__(self, in_feature):
        super().__init__()
        self.l0 = nn.Linear(in_feature, 1024)
        self.l1 = nn.Linear(1024, 512)
        self.l2 = nn.Linear(512, 1)
        # self.apply(init_weights)

    def forward(self, visual, language):

        x = torch.cat((visual, language), dim=-1)
        out = F.relu(self.l0(x))
        out = F.relu(self.l1(out))
        out = self.l2(out)
        return out

class MI(nn.Module):

    @ex.capture
    def __init__(self, visual_size, language_size):
        super(MI,self).__init__()

        self.global_D = GlobalDiscriminator(visual_size+language_size)
        self.ln = nn.LayerNorm([visual_size],elementwise_affine=False)

    def forward(self, visual, language):

        visual = self.ln(visual)
        visual_fake = torch.cat((visual[1:], visual[0:1]), dim=0)
        Ej = -F.softplus(-self.global_D(visual, language)).mean()
        Em = F.softplus(self.global_D(visual_fake, language)).mean()
        dim = Ej - Em

        return dim

    def temp_constrain(self, feat0, feat1, language):

        vis_0 = self.ln(feat0)
        vis_1 = self.ln(feat1)
        visual_fake = torch.cat((vis_0[1:], vis_0[0:1]), dim=0)
        Ej = -F.softplus(-self.global_D(vis_1, language)).mean()
        Em = F.softplus(self.global_D(visual_fake, language)).mean()
        dim = Ej - Em

        return dim

    @ex.capture
    def get_acc(self, visual, unseen_language, label, unseen_label):

        bs = visual.shape[0]
        unsee_num = unseen_language.shape[0]
        visual = repeat(visual,'b c -> b u c',u = unsee_num)
        unseen_language = repeat(unseen_language,'u c -> b u c', b = bs)
        dim_list = -F.softplus(-self.global_D(visual, unseen_language)).squeeze(-1)
        _, pred = torch.max(dim_list, 1)
        unseen_label = torch.tensor(unseen_label).cuda()
        pred = torch.index_select(unseen_label,0,pred)
        acc = pred.eq(label.view_as(pred)).float().mean()

        return acc, pred

@ex.capture
def temp_mask(data, mask_frame):
    x = data.clone()
    n, c, t, v, m = x.shape
    remain_num = t - mask_frame
    remain_frame = random.sample(range(t), remain_num)
    remain_frame.sort()
    x = x[:, :, remain_frame, :, :]

    return x

def motion_att_temp_mask(data, mask_frame):

    n, c, t, v, m = data.shape
    temp = data.clone()
    remain_num = t - mask_frame

    ## get the motion_attention value
    motion = torch.zeros_like(temp)
    motion[:, :, :-1, :, :] = temp[:, :, 1:, :, :] - temp[:, :, :-1, :, :]
    motion = -(motion)**2
    temporal_att = motion.mean((1,3,4))

    ## The frames with the smallest att are reserved
    _,temp_list = torch.topk(temporal_att, remain_num)
    temp_list,_ = torch.sort(temp_list.squeeze())
    temp_list = repeat(temp_list,'n t -> n c t v m',c=c,v=v,m=m)
    temp_resample = temp.gather(2,temp_list)

    ## random temp mask
    random_frame = random.sample(range(remain_num), remain_num-mask_frame)
    random_frame.sort()
    output = temp_resample[:, :, random_frame, :, :]

    return output

def motion_att_temp_mask2(data, mask_frame):

    n, c, t, v, m = data.shape
    temp = data.clone()
    remain_num = t - mask_frame

    ## get the motion_attention value
    motion_pre = torch.zeros_like(temp)
    motion_nex = torch.zeros_like(temp)
    motion_pre[:, :, :-1, :, :] = temp[:, :, 1:, :, :] - temp[:, :, :-1, :, :]
    motion_nex[:, :, 1:, :, :] = temp[:, :, :-1, :, :] - temp[:, :, 1:, :, :]
    motion = -((motion_pre)**2+(motion_nex)**2)
    temporal_att = motion.mean((1,3,4))

    ## The frames with the smallest att are reserved
    _,temp_list = torch.topk(temporal_att, remain_num)
    temp_list,_ = torch.sort(temp_list.squeeze())
    temp_list = repeat(temp_list,'n t -> n c t v m',c=c,v=v,m=m)
    temp_resample = temp.gather(2,temp_list)

    ## random temp mask
    random_frame = random.sample(range(remain_num), remain_num-mask_frame)
    random_frame.sort()
    output = temp_resample[:, :, random_frame, :, :]

    return output