-
Notifications
You must be signed in to change notification settings - Fork 3
/
modules.py
117 lines (85 loc) · 4.69 KB
/
modules.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
import torch
from torch import nn
from torch.autograd import Variable
from torch.nn import functional as tf
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.autograd.set_detect_anomaly(True)
def init_hidden(x, hidden_size: int):
return torch.zeros(1, x.size(0), hidden_size).to(device)
class FeatureLayer(nn.Module):
def __init__(self,input_size):
super(FeatureLayer,self).__init__()
self.input_size = input_size
self.attn_layer = nn.Sequential(nn.Linear(self.input_size, self.input_size),
nn.Tanh(),
nn.Linear(self.input_size, self.input_size))
def forward(self, history_feature, future_feature):
input_data = torch.cat((history_feature, future_feature),1)
attn_weights = tf.softmax(self.attn_layer(input_data), 2)
weighted_input = torch.mul(attn_weights, input_data)
weighted_history_input=weighted_input[:,:history_feature.shape[1],:]
weighted_future_input = weighted_input[:, history_feature.shape[1]:, :]
return attn_weights, weighted_history_input, weighted_future_input
class Encoder(nn.Module):
def __init__(self, input_size: int, hidden_size: int, T: int):
super(Encoder, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.T = T
self.period = int(self.T/24)
self.num_layers = 1
self.lstm_layer = nn.LSTM(input_size=self.input_size, hidden_size=hidden_size, num_layers=self.num_layers, bidirectional=True)
def forward(self, input_data, y_history):
self.lstm_layer.flatten_parameters()
input_data = torch.cat((input_data, y_history), dim=2)
input_encoded, lstm_states = self.lstm_layer(input_data.permute(1,0,2))
hidden = lstm_states[0]
cell = lstm_states[1]
return input_encoded.permute(1,0,2), hidden, cell
class Decoder(nn.Module):
def __init__(self, hidden_size: int, T: int, f_T:int, m_features, out_feats=1):
super(Decoder, self).__init__()
self.input_size = m_features
self.T = T
self.f_T = f_T
self.period = int(self.T/24)
self.encoder_hidden_size = hidden_size
self.decoder_hidden_size = hidden_size
self.attn_layer = nn.Sequential(nn.Linear(2*hidden_size + self.input_size, T),
nn.Tanh(),
nn.Linear(T, T))
self.lstm_layer = nn.LSTM(input_size=2*hidden_size + self.input_size, hidden_size=hidden_size, bidirectional=True, num_layers=1)
self.fc = nn.Linear(hidden_size + out_feats, out_feats)
self.fc_final = nn.Sequential(nn.Linear(2*hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, out_feats))
def forward(self, input_encoded, feature_history1, y_history, feature_future1,hidden,cell):
self.lstm_layer.flatten_parameters()
feature_history = feature_history1.view(feature_history1.shape[0],self.period,-1,feature_history1.shape[2]).permute(0,1,3,2)
feature_future = feature_future1.permute(0,2,1).unsqueeze(1)
hidden_initial = hidden
cell_initial = cell
input_encoded = input_encoded.to(device)
output_final = torch.zeros(input_encoded.size(0), self.f_T, 1).to(device)
feature_future2 = torch.cat([feature_future] * self.period, 1)
cos = nn.PairwiseDistance(p=2)
output = cos(feature_future2.permute(0,3,1,2), feature_history.permute(0,3,1,2))
dp = torch.sum(output, dim=2)
dp = torch.pow(torch.rsqrt(dp),2)
alpha = tf.softmax(dp, dim=1)
for t in range(self.f_T):
x = torch.cat((hidden_initial[0].unsqueeze(0),hidden_initial[1].unsqueeze(0), feature_future[:,:,:,t].permute(1,0,2)),dim=2)
x = tf.softmax(self.attn_layer(x),2).squeeze(0)
c = torch.mul(x.unsqueeze(2),input_encoded)
c = c.view(c.shape[0], int(c.shape[1]/24), -1, c.shape[2])
c = c.sum(2)
c = torch.mul(c.permute(0,2,1), alpha.reshape(alpha.shape[0],1,alpha.shape[1]))
c = c.sum(dim=2)
input_decoder = torch.cat((c, feature_future[:,:,:,t].squeeze(1)), dim=1)
self.lstm_layer.flatten_parameters()
out, lstm_output = self.lstm_layer(input_decoder.unsqueeze(0), (hidden_initial, cell_initial))
hidden_initial = lstm_output[0]
cell_initial = lstm_output[1]
output_final1 = self.fc_final(torch.cat((hidden_initial[0],hidden_initial[1]),dim=1)).squeeze(0)
output_final[:, t, :] = output_final1
return output_final