-
-
Notifications
You must be signed in to change notification settings - Fork 90
/
core.py
90 lines (75 loc) · 2.37 KB
/
core.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
"""
"""
# Created by Wenjie Du <[email protected]>
# License: BSD-3-Clause
import torch.nn as nn
from .submodules import (
SeasonalLayerNorm,
AutoformerEncoderLayer,
AutoCorrelation,
AutoCorrelationLayer,
)
from ...informer.modules.submodules import InformerEncoder
from ....nn.modules.transformer.embedding import DataEmbedding
from ....utils.metrics import calc_mse
class _Autoformer(nn.Module):
def __init__(
self,
n_steps,
n_features,
n_layers,
n_heads,
d_model,
d_ffn,
factor,
moving_avg_window_size,
dropout,
activation="relu",
output_attention=False,
):
super().__init__()
self.seq_len = n_steps
self.n_layers = n_layers
self.enc_embedding = DataEmbedding(
n_features,
d_model,
dropout=dropout,
with_pos=False,
)
self.encoder = InformerEncoder(
[
AutoformerEncoderLayer(
AutoCorrelationLayer(
AutoCorrelation(False, factor, dropout, output_attention),
d_model,
n_heads,
),
d_model,
d_ffn,
moving_avg_window_size,
dropout,
activation,
)
for _ in range(n_layers)
],
norm_layer=SeasonalLayerNorm(d_model),
)
# for the imputation task, the output dim is the same as input dim
self.projection = nn.Linear(d_model, n_features)
def forward(self, inputs: dict, training: bool = True) -> dict:
X, masks = inputs["X"], inputs["missing_mask"]
# embedding
enc_out = self.enc_embedding(X) # [B,T,C]
# Autoformer encoder processing
enc_out, attns = self.encoder(enc_out)
# project back the original data space
dec_out = self.projection(enc_out)
imputed_data = masks * X + (1 - masks) * dec_out
results = {
"imputed_data": imputed_data,
}
if training:
# `loss` is always the item for backward propagating to update the model
loss = calc_mse(dec_out, inputs["X_ori"], inputs["indicating_mask"])
results["loss"] = loss
return results