-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathmodels.py
131 lines (102 loc) · 4.36 KB
/
models.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
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
import torch
import torch.nn as nn
from STConvLSTM import STConvLSTMCell
class Generator(nn.Module):
"""
Generator model with Spatio-Temporal Convolutional LSTMs.
"""
def __init__(self, cfg, device):
super(Generator, self).__init__()
self.input_size = cfg["input_size"]
self.hidden_dim = cfg["hidden_dim"]
self.input_dim = cfg["input_dim"]
self.kernel_size = tuple(cfg["kernel_size"])
self.height, self.width = self.input_size
self.device = device
self.STConvLSTM_Cell_1 = STConvLSTMCell(
input_size=self.input_size,
input_dim=self.input_dim,
hidden_dim=self.hidden_dim,
kernel_size=self.kernel_size,
bias=True,
)
self.STConvLSTM_Cell_2 = STConvLSTMCell(
input_size=self.input_size,
input_dim=self.hidden_dim,
hidden_dim=self.hidden_dim,
kernel_size=self.kernel_size,
bias=True,
)
self.STConvLSTM_Cell_3 = STConvLSTMCell(
input_size=self.input_size,
input_dim=self.hidden_dim,
hidden_dim=self.hidden_dim,
kernel_size=self.kernel_size,
bias=True,
)
self.STConvLSTM_Cell_4 = STConvLSTMCell(
input_size=self.input_size,
input_dim=self.hidden_dim,
hidden_dim=self.hidden_dim,
kernel_size=self.kernel_size,
bias=True,
)
self.head = nn.Conv2d(in_channels=self.hidden_dim, out_channels=self.input_dim, kernel_size=(1, 1))
def forward(self, input_sequence, future=10):
batch_size = input_sequence.size(0)
hidden_initializer = [torch.zeros(batch_size, self.hidden_dim, self.height, self.width).to(self.device)] * 3
h_t1, c_t1, m_t1 = hidden_initializer.copy()
h_t2, c_t2, _ = hidden_initializer.copy()
h_t3, c_t3, _ = hidden_initializer.copy()
h_t4, c_t4, _ = hidden_initializer.copy()
outputs = []
seq_len = input_sequence.size(1)
for time in range(seq_len):
if time:
m_t1 = m_t4
h_t1, c_t1, m_t1 = self.STConvLSTM_Cell_1(
input_tensor=input_sequence[:, time, :, :, :], cur_state=[h_t1, c_t1, m_t1]
)
h_t2, c_t2, m_t2 = self.STConvLSTM_Cell_2(input_tensor=h_t1, cur_state=[h_t2, c_t2, m_t1])
h_t3, c_t3, m_t3 = self.STConvLSTM_Cell_3(input_tensor=h_t2, cur_state=[h_t3, c_t3, m_t2])
h_t4, c_t4, m_t4 = self.STConvLSTM_Cell_4(input_tensor=h_t3, cur_state=[h_t4, c_t4, m_t3])
output = self.head(h_t4)
output = torch.sigmoid(output)
outputs += [output]
for t in range(future):
m_t1 = m_t4
h_t1, c_t1, m_t1 = self.STConvLSTM_Cell_1(input_tensor=outputs[-1], cur_state=[h_t1, c_t1, m_t1])
h_t2, c_t2, m_t2 = self.STConvLSTM_Cell_2(input_tensor=h_t1, cur_state=[h_t2, c_t2, m_t1])
h_t3, c_t3, m_t3 = self.STConvLSTM_Cell_3(input_tensor=h_t2, cur_state=[h_t3, c_t3, m_t2])
h_t4, c_t4, m_t4 = self.STConvLSTM_Cell_4(input_tensor=h_t3, cur_state=[h_t4, c_t4, m_t3])
output = self.head(h_t4)
output = torch.sigmoid(output)
outputs += [output]
outputs = torch.stack(outputs, 1)
return outputs
class Discriminator(nn.Module):
"""
Discriminator model.
"""
def __init__(self, cfg):
super(Discriminator, self).__init__()
self.input_size = cfg["input_size"]
self.hidden_dim = cfg["hidden_dim"]
self.height, self.width = self.input_size
self.linear_1 = nn.Linear(self.height * self.width, self.hidden_dim * 4)
self.linear_2 = nn.Linear(self.hidden_dim * 4, self.hidden_dim * 2)
self.linear_3 = nn.Linear(self.hidden_dim * 2, self.hidden_dim)
self.linear_4 = nn.Linear(self.hidden_dim, 1)
self.dropout = nn.Dropout(0.3)
self.relu = nn.LeakyReLU(0.2, inplace=True)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
x = x.view(-1, self.height * self.width)
x = self.relu(self.linear_1(x))
x = self.dropout(x)
x = self.relu(self.linear_2(x))
x = self.dropout(x)
x = self.relu(self.linear_3(x))
x = self.dropout(x)
out = self.sigmoid(self.linear_4(x))
return out