MultiGradientSensing/model.py at main · MinKim14/MultiGradientSensing · 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
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter


class SimpleLinearModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleLinearModel, self).__init__()
        self.linear = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size),
            nn.Softmax(),
        )

    def forward(self, x):
        out = self.linear(x)
        return out


# LSTM model
class SimpleLSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=3):
        super(SimpleLSTMModel, self).__init__()
        self.lstm = nn.LSTM(
            input_size,
            hidden_size,
            batch_first=True,
            dropout=0.3,
            num_layers=num_layers,
        )
        self.linear = nn.Sequential(
            nn.Linear(hidden_size + input_size, hidden_size),
            nn.Dropout(0.3),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.Dropout(0.3),
            nn.SiLU(),
            nn.Linear(hidden_size, output_size),
            # nn.Softmax(),
        )

    def forward(self, x):
        emb, _ = self.lstm(x)

        emb = torch.cat([emb, x], dim=2)
        out = self.linear(emb[:, -1, :])
        # return out
        return out, emb[:, -1, :]


class SimpleGRUModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=3):
        super(SimpleGRUModel, self).__init__()
        self.lstm = nn.GRU(
            input_size,
            hidden_size,
            batch_first=True,
            dropout=0.1,
            num_layers=num_layers,
        )
        self.linear = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Dropout(0.1),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, output_size),
            nn.Softmax(),
        )

    def forward(self, x):
        emb, _ = self.lstm(x)
        out = self.linear(emb[:, -1, :])
        return out, emb[:, -1, :]


class SimpleLinearModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=3):
        super(SimpleLinearModel, self).__init__()

        self.first_linear = nn.ModuleList()
        self.first_linear.add_module(
            "linear_{}".format(0),
            nn.Sequential(
                nn.Linear(input_size * 32, hidden_size),
                nn.Dropout(0.1),
                nn.ReLU(),
            ),
        )
        for i in range(num_layers - 2):
            self.first_linear.add_module(
                "linear_{}".format(i + 1),
                nn.Sequential(
                    nn.Linear(hidden_size, hidden_size),
                    nn.Dropout(0.1),
                    nn.ReLU(),
                ),
            )

        self.first_linear.add_module(
            "linear_{}".format(num_layers - 1),
            nn.Sequential(
                nn.Linear(hidden_size, output_size),
                nn.Softmax(),
            ),
        )

    def forward(self, x):
        x = x.reshape(x.shape[0], -1)
        for i, layer in enumerate(self.first_linear):
            x = layer(x)
        return x, x