Merge pull request #348 from WenjieDu/(refactor)modular_multitask_nn

Modularize neural network models

docs/index.rst

@@ -170,7 +170,7 @@ Additionally, we present you a usage example of imputing missing values in time
 ❖ Available Algorithms
 ^^^^^^^^^^^^^^^^^^^^^^^
 PyPOTS supports imputation, classification, clustering, and forecasting tasks on multivariate time series with missing values.
-The currently available algorithms of four tasks are cataloged in the following table with four partitions. The paper references are all listed `on the reference page </references.html>`_.
+The currently available algorithms of four tasks are cataloged in the following table with four partitions. The paper references are provided and you can click them to check out.
 
 🌟 Since **v0.2**, all neural-network models in PyPOTS has got hyperparameter-optimization support.
 This functionality is implemented with the `Microsoft NNI <https://github.com/microsoft/nni>`_ framework.
@@ -183,7 +183,7 @@ the same as we did in `SAITS paper <https://arxiv.org/pdf/2202.08516)>`_.
 Task Type Algorithm Year Reference
 ============================== ================ ========================================================================================================= ====== =========
 Imputation Neural Net SAITS (Self-Attention-based Imputation for Time Series) 2023 :cite:`du2023SAITS`
-Imputation Neural Net Transformer 2017 :cite:`vaswani2017Transformer`, :cite:`du2023SAITS`
+Imputation Neural Net Transformer 2017 :cite:`vaswani2017Transformer`
 Imputation Neural Net Crossformer (Transformer Utilizing Cross-Dimension Dependency for Multivariate Time Series Forecasting) 2023 :cite:`nie2023patchtst`
 Imputation Neural Net TimesNet (Temporal 2D-Variation Modeling for General Time Series Analysis) 2023 :cite:`wu2023timesnet`
 Imputation Neural Net PatchTST (A Time Series is Worth 64 Words: Long-Term Forecasting with Transformers) 2023 :cite:`nie2023patchtst`

docs/pypots.nn.rst

@@ -8,14 +8,134 @@ pypots.nn.functional
  :members:
 
 
-pypots.nn.modules.rnn
----------------------
+pypots.nn.modules.autoformer
+-----------------------------
+
+.. automodule:: pypots.nn.modules.autoformer
+ :members:
+
+
+pypots.nn.modules.brits
+-----------------------------
+
+.. automodule:: pypots.nn.modules.brits
+ :members:
+
+
+pypots.nn.modules.crli
+-----------------------------
+
+.. automodule:: pypots.nn.modules.crli
+ :members:
+
+
+pypots.nn.modules.crossformer
+-----------------------------
+
+.. automodule:: pypots.nn.modules.crossformer
+ :members:
+
+
+pypots.nn.modules.csdi
+-----------------------------
+
+.. automodule:: pypots.nn.modules.csdi
+ :members:
+
+
+pypots.nn.modules.dlinear
+-----------------------------
+
+.. automodule:: pypots.nn.modules.dlinear
+ :members:
+
+
+pypots.nn.modules.etsformer
+-----------------------------
+
+.. automodule:: pypots.nn.modules.etsformer
+ :members:
 
-.. automodule:: pypots.nn.modules.rnn
+
+pypots.nn.modules.fedformer
+-----------------------------
+
+.. automodule:: pypots.nn.modules.fedformer
  :members:
 
+
+pypots.nn.modules.gpvae
+-----------------------------
+
+.. automodule:: pypots.nn.modules.gpvae
+ :members:
+
+
+pypots.nn.modules.grud
+-----------------------------
+
+.. automodule:: pypots.nn.modules.grud
+ :members:
+
+
+pypots.nn.modules.informer
+-----------------------------
+
+.. automodule:: pypots.nn.modules.informer
+ :members:
+
+
+pypots.nn.modules.mrnn
+-----------------------------
+
+.. automodule:: pypots.nn.modules.mrnn
+ :members:
+
+
+pypots.nn.modules.patchtst
+-----------------------------
+
+.. automodule:: pypots.nn.modules.patchtst
+ :members:
+
+
+pypots.nn.modules.raindrop
+-----------------------------
+
+.. automodule:: pypots.nn.modules.raindrop
+ :members:
+
+
+pypots.nn.modules.saits
+-----------------------------
+
+.. automodule:: pypots.nn.modules.saits
+ :members:
+
+
+pypots.nn.modules.timesnet
+-----------------------------
+
+.. automodule:: pypots.nn.modules.timesnet
+ :members:
+
+
 pypots.nn.modules.transformer
 -----------------------------
 
 .. automodule:: pypots.nn.modules.transformer
  :members:
+
+
+pypots.nn.modules.usgan
+-----------------------------
+
+.. automodule:: pypots.nn.modules.usgan
+ :members:
+
+
+pypots.nn.modules.vader
+-----------------------------
+
+.. automodule:: pypots.nn.modules.vader
+ :members:

pypots/classification/brits/core.py

@@ -0,0 +1,89 @@
+"""
+The implementation of BRITS for the partially-observed time-series classification task.
+
+Refer to the paper "Cao, W., Wang, D., Li, J., Zhou, H., Li, L., & Li, Y. (2018).
+BRITS: Bidirectional Recurrent Imputation for Time Series. NeurIPS 2018."
+
+Notes
+-----
+Partial implementation uses code from https://github.com/caow13/BRITS. The bugs in the original implementation
+are fixed here.
+
+"""
+
+# Created by Wenjie Du <wenjay.du@gmail.com>
+# License: BSD-3-Clause
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...nn.modules.brits import BackboneBRITS
+
+
+class _BRITS(nn.Module):
+ def __init__(
+ self,
+ n_steps: int,
+ n_features: int,
+ rnn_hidden_size: int,
+ n_classes: int,
+ classification_weight: float,
+ reconstruction_weight: float,
+ ):
+ super().__init__()
+ self.n_steps = n_steps
+ self.n_features = n_features
+ self.rnn_hidden_size = rnn_hidden_size
+ self.n_classes = n_classes
+ self.classification_weight = classification_weight
+ self.reconstruction_weight = reconstruction_weight
+
+ # create models
+ self.model = BackboneBRITS(n_steps, n_features, rnn_hidden_size)
+ self.f_classifier = nn.Linear(self.rnn_hidden_size, n_classes)
+ self.b_classifier = nn.Linear(self.rnn_hidden_size, n_classes)
+
+ def forward(self, inputs: dict, training: bool = True) -> dict:
+ (
+ imputed_data,
+ f_reconstruction,
+ b_reconstruction,
+ f_hidden_states,
+ b_hidden_states,
+ consistency_loss,
+ reconstruction_loss,
+ ) = self.model(inputs)
+
+ f_logits = self.f_classifier(f_hidden_states)
+ b_logits = self.b_classifier(b_hidden_states)
+ f_prediction = torch.softmax(f_logits, dim=1)
+ b_prediction = torch.softmax(b_logits, dim=1)
+ classification_pred = (f_prediction + b_prediction) / 2
+
+ results = {
+ "imputed_data": imputed_data,
+ "classification_pred": classification_pred,
+ }
+
+ # if in training mode, return results with losses
+ if training:
+ results["consistency_loss"] = consistency_loss
+ results["reconstruction_loss"] = reconstruction_loss
+ f_classification_loss = F.nll_loss(torch.log(f_prediction), inputs["label"])
+ b_classification_loss = F.nll_loss(torch.log(b_prediction), inputs["label"])
+ classification_loss = (f_classification_loss + b_classification_loss) / 2
+ loss = (
+ consistency_loss
+ + reconstruction_loss * self.reconstruction_weight
+ + classification_loss * self.classification_weight
+ )
+
+ # `loss` is always the item for backward propagating to update the model
+ results["loss"] = loss
+ results["reconstruction"] = (f_reconstruction + b_reconstruction) / 2
+ results["classification_loss"] = classification_loss
+ results["f_reconstruction"] = f_reconstruction
+ results["b_reconstruction"] = b_reconstruction
+
+ return results

pypots/classification/brits/model.py

@@ -21,8 +21,8 @@
 import torch
 from torch.utils.data import DataLoader
 
+from .core import _BRITS
 from .data import DatasetForBRITS
-from .modules import _BRITS
 from ..base import BaseNNClassifier
 from ...optim.adam import Adam
 from ...optim.base import Optimizer
@@ -135,7 +135,6 @@ def __init__(
  self.n_classes,
  self.classification_weight,
  self.reconstruction_weight,
- self.device,
  )
  self._send_model_to_given_device()
  self._print_model_size()