| 9 | import matplotlib.pyplot as plt |
| 10 | |
| 11 | class Sequence(nn.Module): |
| 12 | def __init__(self): |
| 13 | super(Sequence, self).__init__() |
| 14 | self.lstm1 = nn.LSTMCell(1, 51) |
| 15 | self.lstm2 = nn.LSTMCell(51, 51) |
| 16 | self.linear = nn.Linear(51, 1) |
| 17 | |
| 18 | def forward(self, input, future = 0): |
| 19 | outputs = [] |
| 20 | h_t = torch.zeros(input.size(0), 51, dtype=torch.double) |
| 21 | c_t = torch.zeros(input.size(0), 51, dtype=torch.double) |
| 22 | h_t2 = torch.zeros(input.size(0), 51, dtype=torch.double) |
| 23 | c_t2 = torch.zeros(input.size(0), 51, dtype=torch.double) |
| 24 | |
| 25 | for input_t in input.split(1, dim=1): |
| 26 | h_t, c_t = self.lstm1(input_t, (h_t, c_t)) |
| 27 | h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2)) |
| 28 | output = self.linear(h_t2) |
| 29 | outputs += [output] |
| 30 | for i in range(future):# if we should predict the future |
| 31 | h_t, c_t = self.lstm1(output, (h_t, c_t)) |
| 32 | h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2)) |
| 33 | output = self.linear(h_t2) |
| 34 | outputs += [output] |
| 35 | outputs = torch.cat(outputs, dim=1) |
| 36 | return outputs |
| 37 | |
| 38 | |
| 39 | if __name__ == '__main__': |