fixed all jacobians and hessian modes, depreceated hessian by obs

kept batched_obs analytic hessian, the single observation one was just slower.

Author: davidsebfischer

batchglm/train/tf/nb_glm/estimator.py

@@ -75,18 +75,9 @@ def map_model(idx, data) -> BasicModelGraph:
         with tf.name_scope("loss"):
             loss = tf.reduce_sum(norm_neg_log_likelihood)
 
-        # TODO: remove this and decide for one implementation
-        if pkg_constants.HESSIAN_MODE == "obs":
-            # Only need iterator that yields single observations for hessian mode obs:
-            singleobs_data = dataset.map(fetch_fn, num_parallel_calls=pkg_constants.TF_NUM_THREADS)
-            singleobs_data = singleobs_data.prefetch(1)
-        else:
-            singleobs_data = None
-
         with tf.name_scope("hessians"):
             hessians = Hessians(
                 batched_data=batched_data,
-                singleobs_data=singleobs_data,
                 sample_indices=sample_indices,
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
@@ -268,7 +259,6 @@ def __init__(
                 # Define the hessian on the batched model for newton-rhapson:
                 batch_hessians = Hessians(
                     batched_data=batch_data,
-                    singleobs_data=None,
                     sample_indices=batch_sample_index,
                     constraints_loc=constraints_loc,
                     constraints_scale=constraints_scale,

batchglm/train/tf/nb_glm/hessians.py

@@ -175,7 +175,6 @@ class Hessians:
     def __init__(
             self,
             batched_data: tf.data.Dataset,
-            singleobs_data: tf.data.Dataset,
             sample_indices: tf.Tensor,
             constraints_loc,
             constraints_scale,
@@ -188,7 +187,6 @@ def __init__(
 
         :param batched_data:
             Dataset iterator over mini-batches of data (used for training) or tf.Tensors of mini-batch.
-        :param singleobs_data: Dataset iterator over single observation batches of data.
         :param sample_indices: Indices of samples to be used.
         :param constraints_loc: Constraints for location model.
             Array with constraints in rows and model parameters in columns.
@@ -218,29 +216,14 @@ def __init__(
             evaluation of the hessian via the tf.hessian function,
             which is done by feature for implementation reasons.
         :param iterator: bool
-            Whether an iterator or a tensor (single yield of an iterator) is given
-            in
+            Whether batched_data is an iterator or a tensor (such as single yield of an iterator).
         """
         if constraints_loc is not None and mode != "tf":
             raise ValueError("closed form hessian does not work if constraints_loc is not None")
         if constraints_scale is not None and mode != "tf":
             raise ValueError("closed form hessian does not work if constraints_scale is not None")
 
-        if mode == "obs":
-            logger.info("Performance warning for hessian mode: " +
-                        "obs_batched is strongly recommended as an alternative to obs.")
-            self.hessian = self.byobs(
-                batched_data=singleobs_data,
-                sample_indices=sample_indices,
-                constraints_loc=constraints_loc,
-                constraints_scale=constraints_scale,
-                model_vars=model_vars,
-                batched=False,
-                iterator=iterator,
-                dtype=dtype
-            )
-            self.neg_hessian = tf.negative(self.hessian)
-        elif mode == "obs_batched":
+        if mode == "obs_batched":
             self.hessian = self.byobs(
                 batched_data=batched_data,
                 sample_indices=sample_indices,
@@ -259,6 +242,7 @@ def __init__(
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
                 model_vars=model_vars,
+                iterator=iterator,
                 dtype=dtype
             )
             self.neg_hessian = tf.negative(self.hessian)
@@ -272,6 +256,7 @@ def __init__(
                 constraints_loc=constraints_loc,
                 constraints_scale=constraints_scale,
                 model_vars=model_vars,
+                iterator=iterator,
                 dtype=dtype
             )
             self.hessian = tf.negative(self.neg_hessian)
@@ -542,6 +527,7 @@ def byfeature(
             constraints_loc,
             constraints_scale,
             model_vars: ModelVars,
+            iterator,
             dtype
     ):
         """
@@ -688,15 +674,21 @@ def _red(prev, cur):
         p_shape_a = model_vars.a.shape[0]
         p_shape_b = model_vars.b.shape[0]
 
-        H = op_utils.map_reduce(
-            last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
-            data=batched_data,
-            map_fn=_map,
-            reduce_fn=_red,
-            parallel_iterations=1,
-        )
-        H = H[0]
-        return H
+        if iterator:
+            H = op_utils.map_reduce(
+                last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
+                data=batched_data,
+                map_fn=_map,
+                reduce_fn=_red,
+                parallel_iterations=1
+            )
+        else:
+            H = _map(
+                idx=sample_indices,
+                data=batched_data
+            )
+
+        return H[0]
 
     def tf_byfeature(
             self,
@@ -705,6 +697,7 @@ def tf_byfeature(
             constraints_loc,
             constraints_scale,
             model_vars: ModelVars,
+            iterator,
             dtype
     ) -> List[tf.Tensor]:
         """
@@ -813,11 +806,18 @@ def _map(idx, data):
         def _red(prev, cur):
             return [tf.add(p, c) for p, c in zip(prev, cur)]
 
-        H = op_utils.map_reduce(
-            last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
-            data=batched_data,
-            map_fn=_map,
-            reduce_fn=_red,
-            parallel_iterations=1,
-        )
+        if iterator:
+            H = op_utils.map_reduce(
+                last_elem=tf.gather(sample_indices, tf.size(sample_indices) - 1),
+                data=batched_data,
+                map_fn=_map,
+                reduce_fn=_red,
+                parallel_iterations=1
+            )
+        else:
+            H = _map(
+                idx=sample_indices,
+                data=batched_data
+            )
+
         return H[0]

batchglm/unit_test/test_nb_glm_hessians.py

@@ -53,7 +53,16 @@ def test_compute_hessians(self):
         design_scale = data_utils.design_matrix(sample_description, formula="~ 1 + condition")
 
         input_data = InputData.new(sim.X, design_loc=design_loc, design_scale=design_scale)
-        print(input_data)
+
+        pkg_constants.HESSIAN_MODE = "tf"
+        self.estimator_tf = estimate(input_data)
+        t0_tf = time.time()
+        # tensorflow computes the negative hessian as the
+        # objective is the negative log-likelihood.
+        self.H_tf = self.estimator_tf.hessians
+        t1_tf = time.time()
+        self.estimator_tf.close_session()
+        self.t_tf = t1_tf - t0_tf
 
         pkg_constants.HESSIAN_MODE = "obs_batched"
         self.estimator_ob = estimate(input_data)
@@ -63,14 +72,6 @@ def test_compute_hessians(self):
         self.estimator_ob.close_session()
         self.t_ob = t1_ob - t0_ob
 
-        pkg_constants.HESSIAN_MODE = "obs"
-        self.estimator_ow = estimate(input_data)
-        t0_ow = time.time()
-        self.H_ow = self.estimator_ow.hessians
-        t1_ow = time.time()
-        self.estimator_ow.close_session()
-        self.t_ow = t1_ow - t0_ow
-
         pkg_constants.HESSIAN_MODE = "feature"
         self.estimator_fw = estimate(input_data)
         t0_fw = time.time()
@@ -79,28 +80,15 @@ def test_compute_hessians(self):
         self.estimator_fw.close_session()
         self.t_fw = t1_fw - t0_fw
 
-        pkg_constants.HESSIAN_MODE = "tf"
-        self.estimator_tf = estimate(input_data)
-        t0_tf = time.time()
-        # tensorflow computes the negative hessian as the
-        # objective is the negative log-likelihood.
-        self.H_tf = -self.estimator_tf.hessians
-        t1_tf = time.time()
-        self.estimator_tf.close_session()
-        self.t_tf = t1_tf - t0_tf
-
         i = 1
         print("\n")
-        print("run time observation-wise analytic solution: ", str(self.t_ow))
         print("run time observation batch-wise analytic solution: ", str(self.t_ob))
         print("run time feature-wise analytic solution: ", str(self.t_fw))
         print("run time feature-wise tensorflow solution: ", str(self.t_tf))
-        print("ratio of analytic feature-wise hessian to analytic observation-wise hessian:")
-        print(self.H_tf.values[i, :, :] / self.H_ow.values[i, :, :])
-        print("ratio of analytic feature-wise hessian to analytic observation batch-wise hessian:")
+        print("ratio of tensorflow feature-wise hessian to analytic observation batch-wise hessian:")
         print(self.H_tf.values[i, :, :] / self.H_ob.values[i, :, :])
-        print("ratio of tensorflow feature-wise hessian to analytic observation-wise hessian:")
-        print(self.H_tf.values[i, :, :] / self.H_ow.values[i, :, :])
+        print("ratio of tensorflow feature-wise hessian to analytic feature-wise hessian:")
+        print(self.H_tf.values[i, :, :] / self.H_fw.values[i, :, :])
 
 
 if __name__ == '__main__':

batchglm/unit_test/test_nb_glm_jacobians.py

@@ -35,7 +35,7 @@ def setUp(self):
     def tearDown(self):
         pass
 
-    def test_compute_hessians(self):
+    def test_compute_jacobians(self):
         logging.getLogger("tensorflow").setLevel(logging.ERROR)
         logging.getLogger("batchglm").setLevel(logging.INFO)