From 4c5b02864fc5be28870c0cb60aebb0b463aafcbf Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Fri, 17 Jan 2020 10:26:42 +0100
Subject: [PATCH] added quantization option to model

---
 pts/model/tempflow/tempflow_estimator.py | 4 ++++
 pts/model/tempflow/tempflow_network.py   | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/pts/model/tempflow/tempflow_estimator.py b/pts/model/tempflow/tempflow_estimator.py
index dda4c63..8a777eb 100644
--- a/pts/model/tempflow/tempflow_estimator.py
+++ b/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)
 
diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py
index 3740b64..bbb9c93 100644
--- a/pts/model/tempflow/tempflow_network.py
+++ b/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))