10679769773

Committed 03 Sep 2024 08:56AM UTC coverage: 83.282% (-0.6%) from 83.927%

Build # 10679769773

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #499 from WenjieDu/(feat)add_timemixer

Implement TimeMixer as an imputation model

Run Details

266 of 416 new or added lines in 8 files covered. (63.94%)

3 existing lines in 2 files now uncovered.

10865 of 13046 relevant lines covered (83.28%)

5.0 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

91.3

/pypots/imputation/timemixer/core.py

"""

"""

# Created by Wenjie Du <wenjay.du@gmail.com>
# License: BSD-3-Clause

import torch.nn as nn

from ...nn.functional import (
    nonstationary_norm,
    nonstationary_denorm,
)
from ...nn.modules.timemixer import BackboneTimeMixer
from ...utils.metrics import calc_mse


class _TimeMixer(nn.Module):
    def __init__(
        self,
        n_layers,
        n_steps,
        n_features,
        d_model,
        d_ffn,
        dropout,
        top_k,
        channel_independence,
        decomp_method,
        moving_avg,
        downsampling_layers,
        downsampling_window,
        apply_nonstationary_norm: bool = False,
    ):
        super().__init__()

        self.apply_nonstationary_norm = apply_nonstationary_norm

        self.model = BackboneTimeMixer(
            task_name="imputation",
            n_steps=n_steps,
            n_features=n_features,
            n_pred_steps=None,
            n_pred_features=n_features,
            n_layers=n_layers,
            d_model=d_model,
            d_ffn=d_ffn,
            dropout=dropout,
            channel_independence=channel_independence,
            decomp_method=decomp_method,
            top_k=top_k,
            moving_avg=moving_avg,
            downsampling_layers=downsampling_layers,
            downsampling_window=downsampling_window,
            downsampling_method="avg",
            use_future_temporal_feature=False,
        )

    def forward(self, inputs: dict, training: bool = True) -> dict:
        X, missing_mask = inputs["X"], inputs["missing_mask"]

        if self.apply_nonstationary_norm:
            # Normalization from Non-stationary Transformer
            X, means, stdev = nonstationary_norm(X, missing_mask)

        # TimesMixer processing
        dec_out = self.model.imputation(X, None)

        if self.apply_nonstationary_norm:
            # De-Normalization from Non-stationary Transformer
            dec_out = nonstationary_denorm(dec_out, means, stdev)

        imputed_data = missing_mask * X + (1 - missing_mask) * 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

1	"""	6✔
2
3	"""
4
5	# Created by Wenjie Du <wenjay.du@gmail.com>
6	# License: BSD-3-Clause
7
8	import torch.nn as nn	6✔
9
10	from ...nn.functional import (	6✔
11	nonstationary_norm,
12	nonstationary_denorm,
13	)
14	from ...nn.modules.timemixer import BackboneTimeMixer	6✔
15	from ...utils.metrics import calc_mse	6✔
16
17
18	class _TimeMixer(nn.Module):	6✔
19	def __init__(	6✔
20	self,
21	n_layers,
22	n_steps,
23	n_features,
24	d_model,
25	d_ffn,
26	dropout,
27	top_k,
28	channel_independence,
29	decomp_method,
30	moving_avg,
31	downsampling_layers,
32	downsampling_window,
33	apply_nonstationary_norm: bool = False,
34	):
35	super().__init__()	6✔
36
37	self.apply_nonstationary_norm = apply_nonstationary_norm	6✔
38
39	self.model = BackboneTimeMixer(	6✔
40	task_name="imputation",
41	n_steps=n_steps,
42	n_features=n_features,
43	n_pred_steps=None,
44	n_pred_features=n_features,
45	n_layers=n_layers,
46	d_model=d_model,
47	d_ffn=d_ffn,
48	dropout=dropout,
49	channel_independence=channel_independence,
50	decomp_method=decomp_method,
51	top_k=top_k,
52	moving_avg=moving_avg,
53	downsampling_layers=downsampling_layers,
54	downsampling_window=downsampling_window,
55	downsampling_method="avg",
56	use_future_temporal_feature=False,
57	)
58
59	def forward(self, inputs: dict, training: bool = True) -> dict:	6✔
60	X, missing_mask = inputs["X"], inputs["missing_mask"]	6✔
61
62	if self.apply_nonstationary_norm:	6✔
63	# Normalization from Non-stationary Transformer
NEW 64	X, means, stdev = nonstationary_norm(X, missing_mask)	×
65
66	# TimesMixer processing
67	dec_out = self.model.imputation(X, None)	6✔
68
69	if self.apply_nonstationary_norm:	6✔
70	# De-Normalization from Non-stationary Transformer
NEW 71	dec_out = nonstationary_denorm(dec_out, means, stdev)	×
72
73	imputed_data = missing_mask * X + (1 - missing_mask) * dec_out	6✔
74	results = {	6✔
75	"imputed_data": imputed_data,
76	}
77
78	if training:	6✔
79	# `loss` is always the item for backward propagating to update the model
80	loss = calc_mse(dec_out, inputs["X_ori"], inputs["indicating_mask"])	6✔
81	results["loss"] = loss	6✔
82
83	return results	6✔

WenjieDu / PyPOTS / 10679769773

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous