diff --git a/pypots/imputation/grud/core.py b/pypots/imputation/grud/core.py index 98f368e0..713259c2 100644 --- a/pypots/imputation/grud/core.py +++ b/pypots/imputation/grud/core.py @@ -33,7 +33,7 @@ def __init__( ) self.output_projection = nn.Linear(rnn_hidden_size, n_features) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: """Forward processing of GRU-D. Parameters @@ -66,7 +66,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.training: results["loss"] = calc_mse(reconstruction, X, missing_mask) return results diff --git a/pypots/imputation/grud/model.py b/pypots/imputation/grud/model.py index 008408a9..b2ea4a0d 100644 --- a/pypots/imputation/grud/model.py +++ b/pypots/imputation/grud/model.py @@ -225,7 +225,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputed_data = results["imputed_data"] imputation_collector.append(imputed_data) diff --git a/pypots/imputation/imputeformer/core.py b/pypots/imputation/imputeformer/core.py index ceb81630..46d49aed 100644 --- a/pypots/imputation/imputeformer/core.py +++ b/pypots/imputation/imputeformer/core.py @@ -92,7 +92,7 @@ def __init__( # apply SAITS loss function to Transformer on the imputation task self.saits_loss_func = SaitsLoss(ORT_weight, MIT_weight) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: x, missing_mask = inputs["X"], inputs["missing_mask"] # x: (batch_size, in_steps, num_nodes) @@ -132,7 +132,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/imputeformer/model.py b/pypots/imputation/imputeformer/model.py index 2c1e229c..334badd9 100644 --- a/pypots/imputation/imputeformer/model.py +++ b/pypots/imputation/imputeformer/model.py @@ -283,7 +283,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputed_data = results["imputed_data"] imputation_collector.append(imputed_data) diff --git a/pypots/imputation/koopa/model.py b/pypots/imputation/koopa/model.py index 76c72c1f..ab780553 100644 --- a/pypots/imputation/koopa/model.py +++ b/pypots/imputation/koopa/model.py @@ -295,7 +295,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/micn/core.py b/pypots/imputation/micn/core.py index 11bfa394..153413e2 100644 --- a/pypots/imputation/micn/core.py +++ b/pypots/imputation/micn/core.py @@ -60,7 +60,7 @@ def __init__( # for the imputation task, the output dim is the same as input dim self.saits_loss_func = SaitsLoss(ORT_weight, MIT_weight) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] seasonal_init, trend_init = self.decomp_multi(X) @@ -82,7 +82,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/micn/model.py b/pypots/imputation/micn/model.py index 3456d539..59061d93 100644 --- a/pypots/imputation/micn/model.py +++ b/pypots/imputation/micn/model.py @@ -276,7 +276,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/moderntcn/core.py b/pypots/imputation/moderntcn/core.py index 3ca8e8f5..4c64d2b8 100644 --- a/pypots/imputation/moderntcn/core.py +++ b/pypots/imputation/moderntcn/core.py @@ -66,7 +66,7 @@ def __init__( individual, ) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] if self.apply_nonstationary_norm: @@ -88,7 +88,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.training: loss = calc_mse(reconstruction, inputs["X_ori"], inputs["indicating_mask"]) results["loss"] = loss diff --git a/pypots/imputation/moderntcn/model.py b/pypots/imputation/moderntcn/model.py index 68ed84ba..145a70d8 100644 --- a/pypots/imputation/moderntcn/model.py +++ b/pypots/imputation/moderntcn/model.py @@ -306,7 +306,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/reformer/core.py b/pypots/imputation/reformer/core.py index ec55c7ad..6c74f01d 100644 --- a/pypots/imputation/reformer/core.py +++ b/pypots/imputation/reformer/core.py @@ -54,7 +54,7 @@ def __init__( 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: + def forward(self, inputs: dict) -> 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 @@ -75,7 +75,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/reformer/model.py b/pypots/imputation/reformer/model.py index 072c6894..c61f6a85 100644 --- a/pypots/imputation/reformer/model.py +++ b/pypots/imputation/reformer/model.py @@ -295,7 +295,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/revinscinet/core.py b/pypots/imputation/revinscinet/core.py index 16d199d3..a3188d5d 100644 --- a/pypots/imputation/revinscinet/core.py +++ b/pypots/imputation/revinscinet/core.py @@ -59,7 +59,7 @@ def __init__( # for the imputation task, the output dim is the same as input dim self.saits_loss_func = SaitsLoss(ORT_weight, MIT_weight) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] X = self.revin(X, missing_mask, mode="norm") @@ -80,7 +80,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/revinscinet/model.py b/pypots/imputation/revinscinet/model.py index 9056f552..4d10967c 100644 --- a/pypots/imputation/revinscinet/model.py +++ b/pypots/imputation/revinscinet/model.py @@ -300,7 +300,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/scinet/core.py b/pypots/imputation/scinet/core.py index 4d2b02a1..1df706c7 100644 --- a/pypots/imputation/scinet/core.py +++ b/pypots/imputation/scinet/core.py @@ -57,7 +57,7 @@ def __init__( # for the imputation task, the output dim is the same as input dim self.saits_loss_func = SaitsLoss(ORT_weight, MIT_weight) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] # WDU: the original SCINet paper isn't proposed for imputation task. Hence the model doesn't take @@ -76,7 +76,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/scinet/model.py b/pypots/imputation/scinet/model.py index 4dd0fa27..29dac999 100644 --- a/pypots/imputation/scinet/model.py +++ b/pypots/imputation/scinet/model.py @@ -302,7 +302,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/stemgnn/core.py b/pypots/imputation/stemgnn/core.py index d8d51efb..650e7bb5 100644 --- a/pypots/imputation/stemgnn/core.py +++ b/pypots/imputation/stemgnn/core.py @@ -48,7 +48,7 @@ def __init__( 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: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] # WDU: the original StemGNN paper isn't proposed for imputation task. Hence the model doesn't take @@ -69,7 +69,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/stemgnn/model.py b/pypots/imputation/stemgnn/model.py index ecee2c80..2f75f4d6 100644 --- a/pypots/imputation/stemgnn/model.py +++ b/pypots/imputation/stemgnn/model.py @@ -276,7 +276,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/tcn/core.py b/pypots/imputation/tcn/core.py index c38390b5..7274af14 100644 --- a/pypots/imputation/tcn/core.py +++ b/pypots/imputation/tcn/core.py @@ -45,7 +45,7 @@ def __init__( 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: + def forward(self, inputs: dict) -> 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 @@ -68,7 +68,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/tcn/model.py b/pypots/imputation/tcn/model.py index 2b33251a..28e987f3 100644 --- a/pypots/imputation/tcn/model.py +++ b/pypots/imputation/tcn/model.py @@ -270,7 +270,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/tefn/core.py b/pypots/imputation/tefn/core.py index f71927a6..ca11825e 100644 --- a/pypots/imputation/tefn/core.py +++ b/pypots/imputation/tefn/core.py @@ -32,7 +32,7 @@ def __init__( n_fod, ) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] if self.apply_nonstationary_norm: @@ -51,7 +51,8 @@ def forward(self, inputs: dict, training: bool = True) -> dict: "imputed_data": imputed_data, } - if training: + # if in training mode, return results with losses + if self.training: # `loss` is always the item for backward propagating to update the model loss = calc_mse(out, inputs["X_ori"], inputs["indicating_mask"]) results["loss"] = loss diff --git a/pypots/imputation/tefn/model.py b/pypots/imputation/tefn/model.py index 2925d8a6..6d55bd7c 100644 --- a/pypots/imputation/tefn/model.py +++ b/pypots/imputation/tefn/model.py @@ -250,7 +250,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/tide/core.py b/pypots/imputation/tide/core.py index e826cbeb..876b3ec4 100644 --- a/pypots/imputation/tide/core.py +++ b/pypots/imputation/tide/core.py @@ -82,7 +82,7 @@ def __init__( # 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: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] # # WDU: the original TiDE paper isn't proposed for imputation task. Hence the model doesn't take @@ -112,7 +112,7 @@ def forward(self, inputs: dict, training: bool = True) -> dict: } # if in training mode, return results with losses - if training: + if self.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 diff --git a/pypots/imputation/tide/model.py b/pypots/imputation/tide/model.py index 7b14a5a6..693e5c5d 100644 --- a/pypots/imputation/tide/model.py +++ b/pypots/imputation/tide/model.py @@ -282,7 +282,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return diff --git a/pypots/imputation/timemixer/core.py b/pypots/imputation/timemixer/core.py index c094d2ff..f988c6fc 100644 --- a/pypots/imputation/timemixer/core.py +++ b/pypots/imputation/timemixer/core.py @@ -56,7 +56,7 @@ def __init__( use_future_temporal_feature=False, ) - def forward(self, inputs: dict, training: bool = True) -> dict: + def forward(self, inputs: dict) -> dict: X, missing_mask = inputs["X"], inputs["missing_mask"] if self.apply_nonstationary_norm: @@ -75,7 +75,8 @@ def forward(self, inputs: dict, training: bool = True) -> dict: "imputed_data": imputed_data, } - if training: + # if in training mode, return results with losses + if self.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 diff --git a/pypots/imputation/timemixer/model.py b/pypots/imputation/timemixer/model.py index 89b24011..7ebf10ca 100644 --- a/pypots/imputation/timemixer/model.py +++ b/pypots/imputation/timemixer/model.py @@ -307,7 +307,7 @@ def predict( with torch.no_grad(): for idx, data in enumerate(test_loader): inputs = self._assemble_input_for_testing(data) - results = self.model.forward(inputs, training=False) + results = self.model.forward(inputs) imputation_collector.append(results["imputed_data"]) # Step 3: output collection and return