added quantization option to model

pts/model/tempflow/tempflow_estimator.py

@@ -55,6 +55,7 @@ class TempFlowEstimator(PTSEstimator):
         hidden_size=100,
         n_hidden=2,
         conditioning_length: int = 200,
+        quantize: bool = False,
 
         scaling: bool = True,
         pick_incomplete: bool = False,
@@ -85,6 +86,7 @@ class TempFlowEstimator(PTSEstimator):
         self.hidden_size = hidden_size
         self.n_hidden = n_hidden
         self.conditioning_length = conditioning_length
+        self.quantize = quantize
 
         self.lags_seq = (
             lags_seq
@@ -173,6 +175,7 @@ class TempFlowEstimator(PTSEstimator):
             hidden_size=self.hidden_size,
             n_hidden=self.n_hidden,
             conditioning_length=self.conditioning_length,
+            quantize=self.quantize,
         ).to(device)
 
     def create_predictor(
@@ -200,6 +203,7 @@ class TempFlowEstimator(PTSEstimator):
             hidden_size=self.hidden_size,
             n_hidden=self.n_hidden,
             conditioning_length=self.conditioning_length,
+            quantize=self.quantize,
             num_parallel_samples=self.num_parallel_samples,
         ).to(device)
 

pts/model/tempflow/tempflow_network.py

@@ -27,6 +27,7 @@ class TempFlowTrainingNetwork(nn.Module):
         n_blocks: int,
         hidden_size: int,
         n_hidden: int,
+        quantize: bool,
         cardinality: List[int] = [1],
         embedding_dimension: int = 1,
         scaling: bool = True,
@@ -61,6 +62,7 @@ class TempFlowTrainingNetwork(nn.Module):
             hidden_size=hidden_size,
             cond_label_size=conditioning_length,
         )
+        self.quantize = quantize
 
         self.distr_output = FlowOutput(
             self.flow, input_size=input_size, cond_size=conditioning_length
@@ -382,6 +384,8 @@ 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:
+            target += torch.rand_like(target)
         likelihoods = -self.flow.log_prob(target, distr_args).unsqueeze(-1)
 
         # assert_shape(likelihoods, (-1, seq_len, 1))