Skip to content

Commit

Permalink
refactor: adjust code and docs of TEFN;
Browse files Browse the repository at this point in the history
  • Loading branch information
WenjieDu committed Sep 8, 2024
1 parent b9cba18 commit eed33c9
Show file tree
Hide file tree
Showing 9 changed files with 118 additions and 125 deletions.
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ The paper references and links are all listed at the bottom of this file.
| **Type** | **Algo** | **IMPU** | **FORE** | **CLAS** | **CLUS** | **ANOD** | **Year - Venue** |
|:--------------|:---------------------------------------------------------------------------------------------------------------------------------|:--------:|:--------:|:--------:|:--------:|:--------:|:-------------------|
| LLM | <a href="https://time-series.ai"><img src="https://time-series.ai/static/figs/robot.svg" width="26px"> Time-Series.AI</a> [^36] |||||| `Later in 2024` |
| Neural Net | TEFN[^39] || | | | | `2024 - arXiv` |
| Neural Net | TEFN🧑‍🔧[^39] || | | | | `2024 - arXiv` |
| Neural Net | TimeMixer[^37] || | | | | `2024 - ICLR` |
| Neural Net | iTransformer🧑‍🔧[^24] || | | | | `2024 - ICLR` |
| Neural Net | ModernTCN[^38] || | | | | `2024 - ICLR` |
Expand Down Expand Up @@ -161,7 +161,7 @@ The paper references and links are all listed at the bottom of this file.
| Neural Net | GRU-D[^4] || || | | `2018 - Sci. Rep.` |
| Neural Net | TCN🧑‍🔧[^35] || | | | | `2018 - arXiv` |
| Neural Net | Transformer🧑‍🔧[^2] || | | | | `2017 - NeurIPS` |
| Naive | Lerp || | | | | |
| Naive | Lerp[^40] || | | | | |
| Naive | LOCF/NOCB || | | | | |
| Naive | Mean || | | | | |
| Naive | Median || | | | | |
Expand Down Expand Up @@ -508,4 +508,4 @@ Time-Series.AI</a>
[^39]: Zhan, T., He, Y., Li, Z., & Deng, Y. (
2024). [Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting](https://arxiv.org/abs/2405.06419).
*arXiv 2024*

[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
5 changes: 3 additions & 2 deletions README_zh.md
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ PyPOTS当前支持多变量POTS数据的插补,预测,分类,聚类以及
| **类型** | **算法** | **插补** | **预测** | **分类** | **聚类** | **异常检测** | **年份 - 刊物** |
|:--------------|:---------------------------------------------------------------------------------------------------------------------------------|:------:|:------:|:------:|:------:|:--------:|:-------------------|
| LLM | <a href="https://time-series.ai"><img src="https://time-series.ai/static/figs/robot.svg" width="26px"> Time-Series.AI</a> [^36] |||||| `Later in 2024` |
| Neural Net | TEFN[^39] || | | | | `2024 - arXiv` |
| Neural Net | TEFN🧑‍🔧[^39] || | | | | `2024 - arXiv` |
| Neural Net | TimeMixer[^37] || | | | | `2024 - ICLR` |
| Neural Net | iTransformer🧑‍🔧[^24] || | | | | `2024 - ICLR` |
| Neural Net | ModernTCN[^38] || | | | | `2024 - ICLR` |
Expand Down Expand Up @@ -142,7 +142,7 @@ PyPOTS当前支持多变量POTS数据的插补,预测,分类,聚类以及
| Neural Net | GRU-D[^4] || || | | `2018 - Sci. Rep.` |
| Neural Net | TCN🧑‍🔧[^35] || | | | | `2018 - arXiv` |
| Neural Net | Transformer🧑‍🔧[^2] || | | | | `2017 - NeurIPS` |
| Naive | Lerp || | | | | |
| Naive | Lerp[^40] || | | | | |
| Naive | LOCF/NOCB || | | | | |
| Naive | Mean || | | | | |
| Naive | Median || | | | | |
Expand Down Expand Up @@ -468,3 +468,4 @@ Time-Series.AI</a>
[^39]: Zhan, T., He, Y., Li, Z., & Deng, Y. (
2024). [Time Evidence Fusion Network: Multi-source View in Long-Term Time Series Forecasting](https://arxiv.org/abs/2405.06419).
*arXiv 2024*
[^40]: [Wikipedia: Linear interpolation](https://en.wikipedia.org/wiki/Linear_interpolation)
2 changes: 2 additions & 0 deletions docs/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,8 @@ The paper references are all listed at the bottom of this readme file.
+----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
| Type | Algorithm | IMPU | FORE | CLAS | CLUS | ANOD | Year - Venue |
+================+===========================================================+======+======+======+======+======+=======================+
| Neural Net | TEFN🧑‍🔧 :cite:`zhan2024tefn` || | | | | ``2024 - arXiv`` |
+----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
| Neural Net | TimeMixer :cite:`wang2024timemixer` || | | | | ``2024 - ICLR`` |
+----------------+-----------------------------------------------------------+------+------+------+------+------+-----------------------+
| Neural Net | iTransformer🧑‍🔧 :cite:`liu2024itransformer` || | | | | ``2024 - ICLR`` |
Expand Down
2 changes: 1 addition & 1 deletion pypots/imputation/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,5 +84,5 @@
"Mean",
"Median",
"Lerp",
"TEFN"
"TEFN",
]
10 changes: 5 additions & 5 deletions pypots/imputation/tefn/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,11 @@

class _TEFN(nn.Module):
def __init__(
self,
n_steps,
n_features,
n_fod,
apply_nonstationary_norm,
self,
n_steps,
n_features,
n_fod,
apply_nonstationary_norm,
):
super().__init__()

Expand Down
12 changes: 6 additions & 6 deletions pypots/imputation/tefn/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,11 @@ class DatasetForTEFN(DatasetForSAITS):
"""Actually TEFN 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,
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)
177 changes: 86 additions & 91 deletions pypots/imputation/tefn/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,83 +23,83 @@

class TEFN(BaseNNImputer):
"""The PyTorch implementation of the TEFN model.
TEFN is originally proposed by Zhan et al. in :cite:`zhan2024tefn`.
Parameters
----------
n_steps :
The number of time steps in the time-series data sample.
n_features :
The number of features in the time-series data sample.
n_fod :
The number of frame of discernment in the TEFN model.
apply_nonstationary_norm :
Whether to apply non-stationary normalization to the input data for TimesNet.
Please refer to :cite:`liu2022nonstationary` for details about non-stationary normalization,
which is not the idea of the original TimesNet paper. Hence, we make it optional
and default not to use here.
batch_size :
The batch size for training and evaluating the model.
epochs :
The number of epochs for training the model.
patience :
The patience for the early-stopping mechanism. Given a positive integer, the training process will be
stopped when the model does not perform better after that number of epochs.
Leaving it default as None will disable the early-stopping.
optimizer :
The optimizer for model training.
If not given, will use a default Adam optimizer.
num_workers :
The number of subprocesses to use for data loading.
`0` means data loading will be in the main process, i.e. there won't be subprocesses.
device :
The device for the model to run on. It can be a string, a :class:`torch.device` object, or a list of them.
If not given, will try to use CUDA devices first (will use the default CUDA device if there are multiple),
then CPUs, considering CUDA and CPU are so far the main devices for people to train ML models.
If given a list of devices, e.g. ['cuda:0', 'cuda:1'], or [torch.device('cuda:0'), torch.device('cuda:1')] ,
the model will be parallely trained on the multiple devices (so far only support parallel training on CUDA
devices). Other devices like Google TPU and Apple Silicon accelerator MPS may be added in the future.
saving_path :
The path for automatically saving model checkpoints and tensorboard files (i.e. loss values recorded during
training into a tensorboard file). Will not save if not given.
model_saving_strategy :
The strategy to save model checkpoints. It has to be one of [None, "best", "better", "all"].
No model will be saved when it is set as None.
The "best" strategy will only automatically save the best model after the training finished.
The "better" strategy will automatically save the model during training whenever the model performs
better than in previous epochs.
The "all" strategy will save every model after each epoch training.
verbose :
Whether to print out the training logs during the training process.
"""
TEFN is originally proposed by Zhan et al. in :cite:`zhan2024tefn`.
Parameters
----------
n_steps :
The number of time steps in the time-series data sample.
n_features :
The number of features in the time-series data sample.
n_fod :
The number of frame of discernment in the TEFN model.
apply_nonstationary_norm :
Whether to apply non-stationary normalization to the input data for TimesNet.
Please refer to :cite:`liu2022nonstationary` for details about non-stationary normalization,
which is not the idea of the original TimesNet paper. Hence, we make it optional
and default not to use here.
batch_size :
The batch size for training and evaluating the model.
epochs :
The number of epochs for training the model.
patience :
The patience for the early-stopping mechanism. Given a positive integer, the training process will be
stopped when the model does not perform better after that number of epochs.
Leaving it default as None will disable the early-stopping.
optimizer :
The optimizer for model training.
If not given, will use a default Adam optimizer.
num_workers :
The number of subprocesses to use for data loading.
`0` means data loading will be in the main process, i.e. there won't be subprocesses.
device :
The device for the model to run on. It can be a string, a :class:`torch.device` object, or a list of them.
If not given, will try to use CUDA devices first (will use the default CUDA device if there are multiple),
then CPUs, considering CUDA and CPU are so far the main devices for people to train ML models.
If given a list of devices, e.g. ['cuda:0', 'cuda:1'], or [torch.device('cuda:0'), torch.device('cuda:1')] ,
the model will be parallely trained on the multiple devices (so far only support parallel training on CUDA
devices). Other devices like Google TPU and Apple Silicon accelerator MPS may be added in the future.
saving_path :
The path for automatically saving model checkpoints and tensorboard files (i.e. loss values recorded during
training into a tensorboard file). Will not save if not given.
model_saving_strategy :
The strategy to save model checkpoints. It has to be one of [None, "best", "better", "all"].
No model will be saved when it is set as None.
The "best" strategy will only automatically save the best model after the training finished.
The "better" strategy will automatically save the model during training whenever the model performs
better than in previous epochs.
The "all" strategy will save every model after each epoch training.
verbose :
Whether to print out the training logs during the training process.
"""

def __init__(
self,
n_steps: int,
n_features: int,
n_fod: int = 2,
apply_nonstationary_norm: bool = True,
batch_size: int = 32,
epochs: int = 100,
patience: int = None,
optimizer: Optional[Optimizer] = Adam(),
num_workers: int = 0,
device: Optional[Union[str, torch.device, list]] = None,
saving_path: str = None,
model_saving_strategy: Optional[str] = "best",
verbose: bool = True,
self,
n_steps: int,
n_features: int,
n_fod: int = 2,
apply_nonstationary_norm: bool = True,
batch_size: int = 32,
epochs: int = 100,
patience: int = None,
optimizer: Optional[Optimizer] = Adam(),
num_workers: int = 0,
device: Optional[Union[str, torch.device, list]] = None,
saving_path: str = None,
model_saving_strategy: Optional[str] = "best",
verbose: bool = True,
):
super().__init__(
batch_size,
Expand All @@ -119,12 +119,7 @@ def __init__(
self.n_fod = n_fod

# set up the model
self.model = _TEFN(
n_steps,
n_features,
n_fod,
self.apply_nonstationary_norm
)
self.model = _TEFN(n_steps, n_features, n_fod, self.apply_nonstationary_norm)
self._send_model_to_given_device()
self._print_model_size()

Expand Down Expand Up @@ -164,10 +159,10 @@ def _assemble_input_for_testing(self, data: list) -> dict:
return inputs

def fit(
self,
train_set: Union[dict, str],
val_set: Optional[Union[dict, str]] = None,
file_type: str = "hdf5",
self,
train_set: Union[dict, str],
val_set: Optional[Union[dict, str]] = None,
file_type: str = "hdf5",
) -> None:
# Step 1: wrap the input data with classes Dataset and DataLoader
training_set = DatasetForTEFN(
Expand Down Expand Up @@ -202,9 +197,9 @@ def fit(
self._auto_save_model_if_necessary(confirm_saving=True)

def predict(
self,
test_set: Union[dict, str],
file_type: str = "hdf5",
self,
test_set: Union[dict, str],
file_type: str = "hdf5",
) -> dict:
"""Make predictions for the input data with the trained model.
Expand Down Expand Up @@ -260,9 +255,9 @@ def predict(
return result_dict

def impute(
self,
test_set: Union[dict, str],
file_type: str = "hdf5",
self,
test_set: Union[dict, str],
file_type: str = "hdf5",
) -> np.ndarray:
"""Impute missing values in the given data with the trained model.
Expand Down
Loading

0 comments on commit eed33c9

Please sign in to comment.