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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156 | def TemporalFusionTransformer(
encoder_len=168,
decoder_len=24,
num_past_vars=6,
num_future_vars=3,
num_static_vars=2,
hidden_dim=64,
num_heads=4,
dropout=0.1,
quantiles=(0.1, 0.5, 0.9),
name="tft",
):
past_in = keras.Input((encoder_len, num_past_vars), name="past_inputs")
future_in = keras.Input((decoder_len, num_future_vars), name="future_inputs")
static_in = keras.Input((num_static_vars,), name="static_inputs")
static_embed = layers.Dense(hidden_dim, name="static_embed")(static_in)
static_context_enrichment = GatedResidualNetwork(
hidden_dim, dropout, name="static_ctx_enrichment"
)(static_embed)
static_context_h = GatedResidualNetwork(hidden_dim, dropout, name="static_ctx_h")(
static_embed
)
static_context_c = GatedResidualNetwork(hidden_dim, dropout, name="static_ctx_c")(
static_embed
)
past_var_embeds = [
layers.Dense(hidden_dim, name=f"past_var{i}_embed")(
layers.Lambda(lambda t, idx=i: t[..., idx : idx + 1])(past_in)
)
for i in range(num_past_vars)
]
past_selected, past_weights = VariableSelectionNetwork(
num_past_vars, hidden_dim, dropout, name="past_vsn"
)(past_var_embeds)
future_var_embeds = [
layers.Dense(hidden_dim, name=f"future_var{i}_embed")(
layers.Lambda(lambda t, idx=i: t[..., idx : idx + 1])(future_in)
)
for i in range(num_future_vars)
]
future_selected, future_weights = VariableSelectionNetwork(
num_future_vars, hidden_dim, dropout, name="future_vsn"
)(future_var_embeds)
encoder_out, state_h, state_c = layers.LSTM(
hidden_dim, return_sequences=True, return_state=True, name="lstm_encoder"
)(past_selected, initial_state=[static_context_h, static_context_c])
decoder_out = layers.LSTM(hidden_dim, return_sequences=True, name="lstm_decoder")(
future_selected, initial_state=[state_h, state_c]
)
lstm_out = layers.Concatenate(axis=1, name="concat_encoder_decoder")(
[encoder_out, decoder_out]
)
lstm_in = layers.Concatenate(axis=1, name="concat_selected")(
[past_selected, future_selected]
)
gated = GatedResidualNetwork(hidden_dim, dropout, name="post_lstm_gate")(lstm_out)
x = layers.LayerNormalization(name="post_lstm_norm")(gated + lstm_in)
enriched = GatedResidualNetwork(hidden_dim, dropout, name="static_enrichment")(
x + static_context_enrichment[:, None, :]
)
total_len = encoder_len + decoder_len
mask = causal_mask(total_len)
attn_out = InterpretableMultiHeadAttention(
hidden_dim, num_heads, dropout, name="attention"
)(enriched, mask=mask)
attn_out = layers.Dropout(dropout)(attn_out)
x = layers.LayerNormalization(name="post_attn_norm")(attn_out + enriched)
x = GatedResidualNetwork(hidden_dim, dropout, name="positionwise_ff")(x)
x = layers.LayerNormalization(name="final_norm")(x + lstm_out)
decoder_span = layers.Lambda(
lambda t, L=encoder_len: t[:, L:, :], name="slice_decoder_span"
)(x)
outputs = layers.Dense(len(quantiles), name="quantile_head")(decoder_span)
return keras.Model([past_in, future_in, static_in], outputs, name=name)
|