diff --git a/pts/model/tempflow/tempflow_estimator.py b/pts/model/tempflow/tempflow_estimator.py index 8a777eb..addf955 100644 --- a/pts/model/tempflow/tempflow_estimator.py +++ b/pts/model/tempflow/tempflow_estimator.py @@ -55,7 +55,7 @@ class TempFlowEstimator(PTSEstimator): hidden_size=100, n_hidden=2, conditioning_length: int = 200, - quantize: bool = False, + dequantize: bool = False, scaling: bool = True, pick_incomplete: bool = False, @@ -86,7 +86,7 @@ class TempFlowEstimator(PTSEstimator): self.hidden_size = hidden_size self.n_hidden = n_hidden self.conditioning_length = conditioning_length - self.quantize = quantize + self.dequantize = dequantize self.lags_seq = ( lags_seq @@ -175,7 +175,7 @@ class TempFlowEstimator(PTSEstimator): hidden_size=self.hidden_size, n_hidden=self.n_hidden, conditioning_length=self.conditioning_length, - quantize=self.quantize, + dequantize=self.dequantize, ).to(device) def create_predictor( @@ -203,7 +203,7 @@ class TempFlowEstimator(PTSEstimator): hidden_size=self.hidden_size, n_hidden=self.n_hidden, conditioning_length=self.conditioning_length, - quantize=self.quantize, + dequantize=self.dequantize, num_parallel_samples=self.num_parallel_samples, ).to(device) diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py index 92bc2af..72f769d 100644 --- a/pts/model/tempflow/tempflow_network.py +++ b/pts/model/tempflow/tempflow_network.py @@ -27,7 +27,7 @@ class TempFlowTrainingNetwork(nn.Module): n_blocks: int, hidden_size: int, n_hidden: int, - quantize: bool, + dequantize: bool, cardinality: List[int] = [1], embedding_dimension: int = 1, scaling: bool = True, @@ -62,7 +62,7 @@ class TempFlowTrainingNetwork(nn.Module): hidden_size=hidden_size, cond_label_size=conditioning_length, ) - self.quantize = quantize + self.dequantize = dequantize self.distr_output = FlowOutput( self.flow, input_size=input_size, cond_size=conditioning_length @@ -384,7 +384,7 @@ class TempFlowTrainingNetwork(nn.Module): # we sum the last axis to have the same shape for all likelihoods # (batch_size, subseq_length, 1) - if self.quantize: + if self.dequantize: target += torch.rand_like(target) likelihoods = -self.flow.log_prob(target, distr_args).unsqueeze(-1)