SF2F_PyTorch/models/fusers.py at main · BAI-Yeqi/SF2F_PyTorch · 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
'''
Implementations of different Fusers
'''


import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import math

try:
    from attention import Attention
except:
    from .attention import Attention


class AttentionFuserV1(nn.Module):
    def __init__(self,
                 dimensions,
                 dim_out,
                 attention_type='general',
                 linear_out_layer=True,
                 ignore_tanh=False):
        '''
        1 Layer Attnetion Fuser
        '''
        nn.Module.__init__(self)
        self.attention = Attention(
            dimensions,
            dim_out,
            attention_type=attention_type,
            linear_out_layer=linear_out_layer,
            ignore_tanh=ignore_tanh)

    def forward(self, embeddings):
        '''
        seq_in: tensor with shape [N, D, L]
        '''
        # Embeddings: [B, N, C]
        embeddings, attention_weights = \
            self.attention(embeddings, embeddings)
        # [B, C, N]
        embeddings = embeddings.transpose(1, 2)
        #print('Transposed Shape:', embeddings.shape)
        embeddings = F.avg_pool1d(
            embeddings, embeddings.size()[2], stride=1)
        embeddings = embeddings.view(embeddings.size()[0], -1, 1, 1)
        #print('Attention Fuser Output Shape:', embeddings.shape)
        return embeddings


class AttentionFuserV2(nn.Module):
    def __init__(self,
                 dimensions,
                 attention_type='general'):
        '''
        2-Layer Attention
        '''
        nn.Module.__init__(self)
        self.attention_1 = Attention(
            dimensions,
            dimensions,
            attention_type=attention_type,
            linear_out_layer=True,
            ignore_tanh=False)
        self.attention_2 = Attention(
            dimensions,
            dimensions,
            attention_type=attention_type,
            linear_out_layer=True,
            ignore_tanh=True)

    def forward(self, embeddings):
        '''
        seq_in: tensor with shape [N, D, L]
        '''
        # Embeddings: [B, N, C]
        embeddings, attention_weights = \
            self.attention_1(embeddings, embeddings)
        embeddings, attention_weights = \
            self.attention_2(embeddings, embeddings)
        # [B, C, N]
        embeddings = embeddings.transpose(1, 2)
        #print('Transposed Shape:', embeddings.shape)
        embeddings = F.avg_pool1d(
            embeddings, embeddings.size()[2], stride=1)
        embeddings = embeddings.view(embeddings.size()[0], -1, 1, 1)
        #print('Attention Fuser Output Shape:', embeddings.shape)
        return embeddings


class AttentionFuserV3(nn.Module):
    def __init__(self,
                 dimensions,
                 attention_type='general'):
        '''
        2-Layer Attention with Residual Connection
        '''
        nn.Module.__init__(self)
        self.attention_1 = Attention(
            dimensions,
            dimensions,
            attention_type=attention_type,
            linear_out_layer=True,
            ignore_tanh=False)
        self.attention_2 = Attention(
            dimensions*2,
            dimensions,
            attention_type=attention_type,
            linear_out_layer=True,
            ignore_tanh=True)

    def forward(self, input):
        '''
        seq_in: tensor with shape [N, D, L]
        '''
        # Embeddings: [B, N, C]
        hidden, attention_weights = \
            self.attention_1(input, input)
        hidden = F.normalize(hidden, dim=2)
        # Residual Connection
        combined = torch.cat((hidden, input), dim=2)
        # 2nd Attention
        embeddings, attention_weights = \
            self.attention_2(combined, combined)
        # [B, C, N]
        embeddings = embeddings.transpose(1, 2)
        #print('Transposed Shape:', embeddings.shape)
        embeddings = F.avg_pool1d(
            embeddings, embeddings.size()[2], stride=1)
        embeddings = embeddings.view(embeddings.size()[0], -1, 1, 1)
        #print('Attention Fuser Output Shape:', embeddings.shape)
        return embeddings