TransCoder/Feature_extractors.py at main · smeshk/TransCoder · 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
import torch
import torch.nn as nn
import torch.nn.functional as F


class FE(nn.Module):
    def __init__(self, mod, NS_model, input_size, d_model):
        """
       The
       :param mod: Transmitter ('trx') or receiver ('rec')
       :param input_size:  Last dimension of the input
       :param d_model: Hidden dimension
        """
        super(FE, self).__init__()
        self.mod = mod
        self.NS_model = NS_model
        self.reserve = 3 + 8
        if self.NS_model == 0:
            self.FC1 = nn.Linear(input_size, d_model, bias=True)
        elif self.NS_model == 1:
            self.FC1 = nn.Linear(input_size, d_model * 3, bias=True)
            self.activation1 = nn.ReLU()
            self.FC2 = nn.Linear(d_model * 3, d_model * 3, bias=True)
            self.activation2 = nn.ReLU()
            self.FC3 = nn.Linear(d_model * 3, d_model, bias=True)
        elif self.NS_model == 2:
            self.FC1 = nn.Linear(input_size, d_model * 3, bias=True)
            self.activation1 = nn.GELU()
            self.FC2 = nn.Linear(d_model * 3, d_model * 3, bias=True)
            self.activation2 = nn.GELU()
            self.FC3 = nn.Linear(d_model * 3, d_model, bias=True)
        elif self.NS_model == 3:
            self.FC1 = nn.Linear(input_size, d_model * 2, bias=True)
            self.activation1 = nn.ReLU()
            self.FC2 = nn.Linear(d_model * 2, d_model * 2, bias=True)
            self.activation2 = nn.ReLU()
            self.FC3 = nn.Linear(d_model * 2, d_model * 2, bias=True)
            self.FC4 = nn.Linear(d_model * 4, d_model, bias=True)
        elif self.NS_model == 4:
            out_size = d_model
            if self.mod == 'trx':
                input_size = input_size - self.reserve
                out_size = d_model - self.reserve
            self.FC1 = nn.Linear(input_size, d_model * 3, bias=True)
            self.activation1 = nn.ReLU()
            self.FC2 = nn.Linear(d_model * 3, d_model * 3, bias=True)
            self.activation2 = nn.ReLU()
            self.FC3 = nn.Linear(d_model * 3, out_size, bias=True)
        elif self.NS_model == 5:
            out_size = d_model * 3
            if self.mod == 'trx':
                input_size = input_size - self.reserve
                out_size = out_size + self.reserve
            self.FC1 = nn.Linear(input_size, d_model * 3, bias=True)
            self.activation1 = nn.ReLU()
            self.FC2 = nn.Linear(d_model * 3, d_model * 3, bias=True)
            self.activation2 = nn.ReLU()
            self.FC3 = nn.Linear(out_size, d_model, bias=True)
        elif self.NS_model == 6:
            self.FC1 = nn.Linear(input_size, d_model * 2, bias=True)
            self.activation1 = nn.ReLU()
            self.FC2 = nn.Linear(d_model * 2, d_model * 2, bias=True)
            self.activation2 = nn.ReLU()
            self.FC3 = nn.Linear(d_model * 2, d_model * 2, bias=True)
            self.FC4 = nn.Linear(d_model * 4, d_model, bias=True)

    def forward(self, src):
        if self.NS_model == 0:
            x = self.FC1(src)
        elif self.NS_model == 1:
            x = self.FC1(src)
            x = self.FC2(self.activation1(x))
            x = self.FC3(self.activation2(x))
        elif self.NS_model == 2:
            x = self.FC1(src)
            x = self.FC2(self.activation1(x))
            x = self.FC3(self.activation2(x))
        elif self.NS_model == 3:
            x1 = self.FC1(src)
            x1 = self.FC2(self.activation1(x1))
            x1 = self.FC3(self.activation2(x1))
            if self.mod == 'trx':  # In the transmitter mode
                constantM = torch.ones_like(src[0])
                constantM[:, self.reserve:] = constantM[:, self.reserve:] * (-1)
                src1 = src * constantM
            elif self.mod == 'rec':
                src1 = src * (-1)
            elif self.mod == 'rec2':
                src1 = src * (-1)

            x2 = self.FC1(src1)  ########## => input_size to d_model dimension
            x2 = self.FC2(self.activation1(x2))
            x2 = self.FC3(self.activation2(x2))
            x = self.FC4(torch.cat([x1, x2], dim=2))
        elif self.NS_model == 4:
            if self.mod == 'trx':  # In the transmitter mode
                src1 = src[:, :, :self.reserve]
                x = self.FC1(src[:, :, self.reserve:])
                x = self.FC2(self.activation1(x))
                x = self.FC3(self.activation2(x))
                x = torch.cat([src1, x], dim=2)
            elif self.mod == 'rec':  # In the decoder mode
                x = self.FC1(src)
                x = self.FC2(self.activation1(x))
                x = self.FC3(self.activation2(x))
            elif self.mod == 'rec2':  # In the belief mode, not used here
                x = self.FC1(src)
                x = self.FC2(self.activation1(x))
                x = self.FC3(self.activation2(x))
        elif self.NS_model == 5:
            if self.mod == 'trx':  # In the transmitter mode
                src1 = src[:, :, :self.reserve]
                x = self.FC1(src[:, :, self.reserve:])
                x = self.FC2(self.activation1(x))
                x = self.activation2(x)
                x = torch.cat([src1, x], dim=2)
                x = self.FC3(x)
            elif self.mod == 'rec':  # In the decoder mode
                x = self.FC1(src)
                x = self.FC2(self.activation1(x))
                x = self.FC3(self.activation2(x))
            elif self.mod == 'rec2':  # In the belief mode, not used here
                x = self.FC1(src)
                x = self.FC2(self.activation1(x))
                x = self.FC3(self.activation2(x))
        elif self.NS_model == 6:
            x1 = self.FC1(src)
            x1 = self.FC2(self.activation1(x1))
            if self.mod == 'trx':  # In the transmitter mode
                constantM = torch.ones_like(src[0])
                constantM[:, self.reserve:] = constantM[:, self.reserve:] * - 1
                src1 = src * constantM
            elif self.mod == 'rec':
                src1 = src * -1
            elif self.mod == 'rec2':  # belief update, not used here
                src1 = src * -1

            x2 = self.FC1(src1)  ########## => input_size to d_model dimension
            x2 = self.FC2(self.activation1(x2))
            x = self.FC4(torch.cat([x1, x2], dim=2))
        return x