uber-time-series/models/encoder_decoder_dropout.py at master · aybchan/uber-time-series · 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
import torch
import torch.nn as nn
from models import variational_dropout as vd

class VDEncoder(nn.Module):
    def __init__(self, in_features, out_features, p):
        super(VDEncoder, self).__init__()
        self.model = nn.ModuleDict({
            'lstm1': vd.LSTM(in_features, 32,dropouto=p, batch_first=True),
            'lstm2': vd.LSTM(32, 8, dropouto=p, batch_first=True),
            'lstm3': vd.LSTM(8, out_features, dropouto=p, batch_first=True),
            'relu':  nn.ReLU()
        })

    def forward(self, x):
        out, _ = self.model['lstm1'](x)
        out, _ = self.model['lstm2'](out)
        out, _ = self.model['lstm3'](out)
        out    = self.model['relu'](out)

        return out


class VDDecoder(nn.Module):
    def __init__(self, p):
        super(VDDecoder, self).__init__()
        self.model = nn.ModuleDict({
            'lstm1': vd.LSTM(1, 2, dropouto=p, batch_first=True),
            'lstm2': vd.LSTM(2, 2, dropouto=p, batch_first=True),
            'lstm3': vd.LSTM(2, 1, dropouto=p, batch_first=True)
        })

    def forward(self, x):
        out, _ = self.model['lstm1'](x)
        out, _ = self.model['lstm2'](out)
        out, _ = self.model['lstm3'](out)

        return out


class VDEncoderDecoder(nn.Module):
    def __init__(self, in_features, input_steps, output_steps, p):
        super(VDEncoderDecoder, self).__init__()
        self.enc_in_features = in_features
        self.input_steps = input_steps # t in the paper
        self.output_steps = output_steps # f in the paper
        self.enc_out_features = 1
        self.traffic_col = 4
        self.p = p

        self.model = nn.ModuleDict({
            'encoder': VDEncoder(self.enc_in_features, self.enc_out_features, self.p),
            'decoder': VDDecoder(self.p),
            'fc1': nn.Linear(self.input_steps + self.output_steps, 32),
            'fc2': nn.Linear(32, self.output_steps)
        })

    def forward(self, x):
        out = self.model['encoder'](x)

        x_auxiliary = x[:,-self.output_steps:,[self.traffic_col]]
        decoder_input = torch.cat([out, x_auxiliary], dim=1)

        out = self.model['decoder'](decoder_input)
        out = self.model['fc1'](out.view(-1, self.input_steps + self.output_steps))
        out = self.model['fc2'](out)

        return out