-
-
Notifications
You must be signed in to change notification settings - Fork 120
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #433 from WenjieDu/(feat)add_reformer
Add Reformer as an imputation model
- Loading branch information
Showing
11 changed files
with
1,756 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
""" | ||
The package of the partially-observed time-series imputation model Reformer. | ||
Refer to the paper | ||
`Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. | ||
Reformer: The Efficient Transformer. | ||
International Conference on Learning Representations, 2020. | ||
<https://openreview.net/pdf?id=rkgNKkHtvB>`_ | ||
Notes | ||
----- | ||
This implementation is inspired by the official one https://github.com/google/trax/tree/master/trax/models/reformer and | ||
https://github.com/lucidrains/reformer-pytorch | ||
""" | ||
|
||
# Created by Wenjie Du <[email protected]> | ||
# License: BSD-3-Clause | ||
|
||
|
||
from .model import Reformer | ||
|
||
__all__ = [ | ||
"Reformer", | ||
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,88 @@ | ||
""" | ||
The core wrapper assembles the submodules of Reformer imputation model | ||
and takes over the forward progress of the algorithm. | ||
""" | ||
|
||
# Created by Wenjie Du <[email protected]> | ||
# License: BSD-3-Clause | ||
|
||
import torch.nn as nn | ||
|
||
from ...nn.modules.reformer import ReformerEncoder | ||
from ...nn.modules.saits import SaitsLoss, SaitsEmbedding | ||
|
||
|
||
class _Reformer(nn.Module): | ||
def __init__( | ||
self, | ||
n_steps, | ||
n_features, | ||
n_layers, | ||
d_model, | ||
n_heads, | ||
bucket_size, | ||
n_hashes, | ||
causal, | ||
d_ffn, | ||
dropout, | ||
ORT_weight: float = 1, | ||
MIT_weight: float = 1, | ||
): | ||
super().__init__() | ||
|
||
self.n_steps = n_steps | ||
|
||
self.saits_embedding = SaitsEmbedding( | ||
n_features * 2, | ||
d_model, | ||
with_pos=False, | ||
dropout=dropout, | ||
) | ||
self.encoder = ReformerEncoder( | ||
n_steps, | ||
n_layers, | ||
d_model, | ||
n_heads, | ||
bucket_size, | ||
n_hashes, | ||
causal, | ||
d_ffn, | ||
dropout, | ||
) | ||
|
||
# for the imputation task, the output dim is the same as input dim | ||
self.output_projection = nn.Linear(d_model, n_features) | ||
self.saits_loss_func = SaitsLoss(ORT_weight, MIT_weight) | ||
|
||
def forward(self, inputs: dict, training: bool = True) -> dict: | ||
X, missing_mask = inputs["X"], inputs["missing_mask"] | ||
|
||
# WDU: the original Reformer paper isn't proposed for imputation task. Hence the model doesn't take | ||
# the missing mask into account, which means, in the process, the model doesn't know which part of | ||
# the input data is missing, and this may hurt the model's imputation performance. Therefore, I apply the | ||
# SAITS embedding method to project the concatenation of features and masks into a hidden space, as well as | ||
# the output layers to project back from the hidden space to the original space. | ||
enc_out = self.saits_embedding(X, missing_mask) | ||
|
||
# Reformer encoder processing | ||
enc_out = self.encoder(enc_out) | ||
# project back the original data space | ||
reconstruction = self.output_projection(enc_out) | ||
|
||
imputed_data = missing_mask * X + (1 - missing_mask) * reconstruction | ||
results = { | ||
"imputed_data": imputed_data, | ||
} | ||
|
||
# if in training mode, return results with losses | ||
if training: | ||
X_ori, indicating_mask = inputs["X_ori"], inputs["indicating_mask"] | ||
loss, ORT_loss, MIT_loss = self.saits_loss_func( | ||
reconstruction, X_ori, missing_mask, indicating_mask | ||
) | ||
results["ORT_loss"] = ORT_loss | ||
results["MIT_loss"] = MIT_loss | ||
# `loss` is always the item for backward propagating to update the model | ||
results["loss"] = loss | ||
|
||
return results |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,24 @@ | ||
""" | ||
Dataset class for Reformer. | ||
""" | ||
|
||
# Created by Wenjie Du <[email protected]> | ||
# License: BSD-3-Clause | ||
|
||
from typing import Union | ||
|
||
from ..saits.data import DatasetForSAITS | ||
|
||
|
||
class DatasetForReformer(DatasetForSAITS): | ||
"""Actually Reformer uses the same data strategy as SAITS, needs MIT for training.""" | ||
|
||
def __init__( | ||
self, | ||
data: Union[dict, str], | ||
return_X_ori: bool, | ||
return_y: bool, | ||
file_type: str = "hdf5", | ||
rate: float = 0.2, | ||
): | ||
super().__init__(data, return_X_ori, return_y, file_type, rate) |
Oops, something went wrong.