Merge pull request #468 from WenjieDu/dev

Add TCN and update docs

README.md

@@ -144,6 +144,7 @@ The paper references and links are all listed at the bottom of this file.
 | Neural Net    | M-RNN[^9]                   |    ✅     |          |          |          |          | `2019 - TBME`      |
 | Neural Net    | BRITS[^3]                   |    ✅     |          |    ✅     |          |          | `2018 - NeurIPS`   |
 | Neural Net    | GRU-D[^4]                   |    ✅     |          |    ✅     |          |          | `2018 - Sci. Rep.` |
+| Neural Net    | TCN🧑‍🔧[^35]               |    ✅     |          |          |          |          | `2018 - arXiv`     |
 | Neural Net    | Transformer🧑‍🔧[^2]        |    ✅     |          |          |          |          | `2017 - NeurIPS`   |
 | Naive         | Lerp                        |    ✅     |          |          |          |          |                    |
 | Naive         | LOCF/NOCB                   |    ✅     |          |          |          |          |                    |
@@ -383,6 +384,7 @@ PyPOTS community is open, transparent, and surely friendly. Let's work together
 [^32]: Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). [Reformer: The Efficient Transformer](https://openreview.net/forum?id=0EXmFzUn5I). *ICLR 2020*.
 [^33]: Cao, D., Wang, Y., Duan, J., Zhang, C., Zhu, X., Huang, C., Tong, Y., Xu, B., Bai, J., Tong, J., & Zhang, Q. (2020). [Spectral Temporal Graph Neural Network for Multivariate Time-series Forecasting](https://proceedings.neurips.cc/paper/2020/hash/cdf6581cb7aca4b7e19ef136c6e601a5-Abstract.html). *NeurIPS 2020*.
 [^34]: Nie, T., Qin, G., Mei, Y., & Sun, J. (2024). [ImputeFormer: Low Rankness-Induced Transformers for Generalizable Spatiotemporal Imputation](https://arxiv.org/abs/2312.01728). *KDD 2024*.
+[^35]: Bai, S., Kolter, J. Z., & Koltun, V. (2018). [An empirical evaluation of generic convolutional and recurrent networks for sequence modeling](https://arxiv.org/abs/1803.01271). *arXiv 2018*.
 
 <details>
 <summary>🏠 Visits</summary>

README_zh.md

@@ -130,6 +130,7 @@ PyPOTS当前支持多变量POTS数据的插补，预测，分类，聚类以及
 | Neural Net    | M-RNN[^9]                   |    ✅     |          |          |          |          | `2019 - TBME`      |
 | Neural Net    | BRITS[^3]                   |    ✅     |          |    ✅     |          |          | `2018 - NeurIPS`   |
 | Neural Net    | GRU-D[^4]                   |    ✅     |          |    ✅     |          |          | `2018 - Sci. Rep.` |
+| Neural Net    | TCN🧑‍🔧[^35]               |    ✅     |          |          |          |          | `2018 - arXiv`     |
 | Neural Net    | Transformer🧑‍🔧[^2]        |    ✅     |          |          |          |          | `2017 - NeurIPS`   |
 | Naive         | Lerp                        |    ✅     |          |          |          |          |                    |
 | Naive         | LOCF/NOCB                   |    ✅     |          |          |          |          |                    |
@@ -354,6 +355,7 @@ PyPOTS社区是一个开放、透明、友好的社区，让我们共同努力
 [^32]: Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). [Reformer: The Efficient Transformer](https://openreview.net/forum?id=0EXmFzUn5I). *ICLR 2020*.
 [^33]: Cao, D., Wang, Y., Duan, J., Zhang, C., Zhu, X., Huang, C., Tong, Y., Xu, B., Bai, J., Tong, J., & Zhang, Q. (2020). [Spectral Temporal Graph Neural Network for Multivariate Time-series Forecasting](https://proceedings.neurips.cc/paper/2020/hash/cdf6581cb7aca4b7e19ef136c6e601a5-Abstract.html). *NeurIPS 2020*.
 [^34]: Nie, T., Qin, G., Mei, Y., & Sun, J. (2024). [ImputeFormer: Low Rankness-Induced Transformers for Generalizable Spatiotemporal Imputation](https://arxiv.org/abs/2312.01728). *KDD 2024*.
+[^35]: Bai, S., Kolter, J. Z., & Koltun, V. (2018). [An empirical evaluation of generic convolutional and recurrent networks for sequence modeling](https://arxiv.org/abs/1803.01271). *arXiv 2018*.
 
 
 <details>

docs/index.rst

@@ -195,6 +195,8 @@ The paper references are all listed at the bottom of this readme file.
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | GRU-D :cite:`che2018GRUD`                                 |      |      |  ✅  |      |      | ``2018 - Sci. Rep.``  |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
+| Neural Net     | TCN🧑‍🔧 :cite:`bai2018tcn`                               |  ✅  |      |      |      |      | ``2018 - arXiv``      |
++----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Neural Net     | Transformer🧑‍🔧 :cite:`vaswani2017Transformer`           |  ✅  |      |      |      |      | ``2017 - NeurIPS``    |
 +----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
 | Naive          | Lerp (Linear Interpolation)                               |  ✅  |      |      |      |      |                       |

docs/pypots.imputation.rst

@@ -261,6 +261,15 @@ pypots.imputation.grud
    :show-inheritance:
    :inherited-members:
 
+pypots.imputation.tcn
+------------------------------
+
+.. automodule:: pypots.imputation.tcn
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :inherited-members:
+
 pypots.imputation.lerp
 -----------------------------
 

docs/references.bib

@@ -756,3 +756,10 @@ @article{nie2024imputeformer
 doi = {10.1145/3637528.3671751},
 url = {https://doi.org/10.1145/3637528.3671751},
 }
+
+@article{bai2018tcn,
+title={An empirical evaluation of generic convolutional and recurrent networks for sequence modeling},
+author={Bai, Shaojie and Kolter, J Zico and Koltun, Vladlen},
+journal={arXiv preprint arXiv:1803.01271},
+year={2018}
+}

pypots/imputation/__init__.py

@@ -23,6 +23,7 @@
 from .crossformer import Crossformer
 from .informer import Informer
 from .autoformer import Autoformer
+from .tcn import TCN
 from .reformer import Reformer
 from .dlinear import DLinear
 from .patchtst import PatchTST
@@ -57,6 +58,7 @@
     "DLinear",
     "Informer",
     "Autoformer",
+    "TCN",
     "Reformer",
     "NonstationaryTransformer",
     "Pyraformer",

pypots/imputation/tcn/__init__.py

@@ -0,0 +1,24 @@
+"""
+The package of the partially-observed time-series imputation model TCN.
+
+Refer to the paper
+`Shaojie Bai, J. Zico Kolter, and Vladlen Koltun.
+"An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling".
+arXiv preprint arXiv:1803.01271.
+<https://arxiv.org/pdf/1803.01271>`_
+
+Notes
+-----
+This implementation is inspired by the official one https://github.com/locuslab/TCN
+
+"""
+
+# Created by Wenjie Du <[email protected]>
+# License: BSD-3-Clause
+
+
+from .model import TCN
+
+__all__ = [
+    "TCN",
+]

pypots/imputation/tcn/core.py

@@ -0,0 +1,81 @@
+"""
+The core wrapper assembles the submodules of TCN 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.saits import SaitsLoss, SaitsEmbedding
+from ...nn.modules.tcn import BackboneTCN
+
+
+class _TCN(nn.Module):
+    def __init__(
+        self,
+        n_steps: int,
+        n_features: int,
+        n_levels: int,
+        d_hidden: int,
+        kernel_size: int,
+        dropout: float = 0,
+        ORT_weight: float = 1,
+        MIT_weight: float = 1,
+    ):
+        super().__init__()
+        self.n_steps = n_steps
+        channel_sizes = [d_hidden] * n_levels
+
+        self.saits_embedding = SaitsEmbedding(
+            n_features * 2,
+            n_features,
+            with_pos=False,
+            dropout=dropout,
+        )
+        self.backbone = BackboneTCN(
+            n_features,
+            channel_sizes,
+            kernel_size,
+            dropout,
+        )
+
+        # for the imputation task, the output dim is the same as input dim
+        self.output_projection = nn.Linear(channel_sizes[-1], 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 TCN 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)
+        enc_out = enc_out.permute(0, 2, 1)
+
+        # TCN encoder processing
+        enc_out = self.backbone(enc_out)
+        enc_out = enc_out.permute(0, 2, 1)
+        # 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

pypots/imputation/tcn/data.py

@@ -0,0 +1,24 @@
+"""
+Dataset class for TCN.
+"""
+
+# Created by Wenjie Du <[email protected]>
+# License: BSD-3-Clause
+
+from typing import Union
+
+from ..saits.data import DatasetForSAITS
+
+
+class DatasetForTCN(DatasetForSAITS):
+    """Actually TCN 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)