Add automatic type casting of input data to nb_glm estimator

Author: Hoeze

batchglm/models/nb/base.py

@@ -47,7 +47,7 @@ def param_shapes(cls) -> dict:
         return INPUT_DATA_PARAMS
 
     @classmethod
-    def new(cls, data, observation_names=None, feature_names=None, dtype="float32"):
+    def new(cls, data, observation_names=None, feature_names=None, cast_dtype=None):
         """
         Create a new InputData object.
 
@@ -62,12 +62,14 @@ def new(cls, data, observation_names=None, feature_names=None, dtype="float32"):
                 stored as data[design_loc] and data[design_scale]
         :param observation_names: (optional) names of the observations.
         :param feature_names: (optional) names of the features.
-        :param dtype: data type of all data; should be either float32 or float64
+        :param cast_dtype: data type of all data; should be either float32 or float64
         :return: InputData object
         """
         X = data_utils.xarray_from_data(data)
-        X = X.astype(dtype)
-        # X = X.chunk({"observations": 1})
+
+        if cast_dtype is not None:
+            X = X.astype(cast_dtype)
+            # X = X.chunk({"observations": 1})
 
         retval = cls(xr.Dataset({
             "X": X,

batchglm/models/nb_glm/base.py

@@ -87,7 +87,7 @@ def new(
             feature_names=None,
             design_loc_key="design_loc",
             design_scale_key="design_scale",
-            dtype="float32"
+            cast_dtype=None
     ):
         """
         Create a new InputData object.
@@ -116,22 +116,23 @@ def new(
         :param feature_names: (optional) names of the features.
         :param design_loc_key: Where to find `design_loc` if `data` is some anndata.AnnData or xarray.Dataset.
         :param design_scale_key: Where to find `design_scale` if `data` is some anndata.AnnData or xarray.Dataset.
-        :param dtype: data type of all data; should be either float32 or float64
+        :param cast_dtype: If this option is set, all provided data will be casted to this data type.
         :return: InputData object
         """
         retval = super(InputData, cls).new(
             data=data,
             observation_names=observation_names,
             feature_names=feature_names,
-            dtype=dtype
+            cast_dtype=cast_dtype
         )
 
         design_loc = _parse_design(data, design_loc, design_loc_names, design_loc_key, "design_loc_params")
         design_scale = _parse_design(data, design_scale, design_scale_names, design_scale_key, "design_scale_params")
 
-        design_loc = design_loc.astype(dtype)
-        design_scale = design_scale.astype(dtype)
-        # design = design.chunk({"observations": 1})
+        if cast_dtype is not None:
+            design_loc = design_loc.astype(cast_dtype)
+            design_scale = design_scale.astype(cast_dtype)
+            # design = design.chunk({"observations": 1})
 
         retval.design_loc = design_loc
         retval.design_scale = design_scale

batchglm/train/tf/nb_glm/estimator.py

@@ -310,7 +310,7 @@ def __init__(
         dataset = tf.data.Dataset.from_tensor_slices(sample_indices)
 
         batched_data = dataset.batch(batch_size)
-        batched_data = batched_data.map(fetch_fn)
+        batched_data = batched_data.map(fetch_fn, num_parallel_calls=pkg_constants.TF_NUM_THREADS)
         batched_data = batched_data.prefetch(1)
 
         def map_model(idx, data) -> BasicModelGraph:
@@ -434,7 +434,7 @@ def __init__(
                 training_data = data_indices.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=2 * batch_size))
                 # training_data = training_data.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
                 training_data = training_data.batch(batch_size, drop_remainder=True)
-                training_data = training_data.map(fetch_fn)
+                training_data = training_data.map(fetch_fn, num_parallel_calls=pkg_constants.TF_NUM_THREADS)
                 training_data = training_data.prefetch(buffer_size)
 
                 iterator = training_data.make_one_shot_iterator()
@@ -768,7 +768,7 @@ class TrainingStrategy(Enum):
                 "stop_at_loss_change": 0.05,
                 "loss_window_size": 10,
                 "use_batching": False,
-                "optim_algo": "GD",
+                "optim_algo": "ADAM",
             },
         ]
         EXACT = [
@@ -814,7 +814,7 @@ class TrainingStrategy(Enum):
                 "stop_at_loss_change": 0.25,
                 "loss_window_size": 10,
                 "use_batching": False,
-                "optim_algo": "GD",
+                "optim_algo": "ADAM",
             },
         ]
 
@@ -837,6 +837,7 @@ def __init__(
             quick_scale=False,
             model: EstimatorGraph = None,
             extended_summary=False,
+            dtype="float64",
     ):
         """
         Create a new Estimator
@@ -871,8 +872,6 @@ def __init__(
         :param extended_summary: Include detailed information in the summaries.
             Will drastically increase runtime of summary writer, use only for debugging.
         """
-        dtype = input_data.X.dtype
-
         # validate design matrix:
         if np.linalg.matrix_rank(input_data.design_loc) != np.linalg.matrix_rank(input_data.design_loc.T):
             raise ValueError("design_loc matrix is not full rank")
@@ -920,9 +919,9 @@ def __init__(
 
                     logger.info("Using closed-form MLE initialization for mean")
                     logger.debug("RMSE of closed-form mean:\n%s", a_prime[1])
-                    logger.debug("Should train mu:\t%s", self._train_mu)
+                    logger.info("Should train mu: %s", self._train_mu)
                 except np.linalg.LinAlgError:
-                    pass
+                    logger.warning("Closed form initialization failed!")
 
             if isinstance(init_b, str) and (init_b.lower() == "auto" or init_b.lower() == "closed_form"):
                 try:
@@ -951,9 +950,9 @@ def __init__(
 
                     logger.info("Using closed-form MME initialization for dispersion")
                     logger.debug("RMSE of closed-form dispersion:\n%s", b_prime[1])
-                    logger.debug("Should train r:\t%s", self._train_r)
+                    logger.info("Should train r: %s", self._train_r)
                 except np.linalg.LinAlgError:
-                    pass
+                    logger.warning("Closed form initialization failed!")
 
             if init_model is not None:
                 if isinstance(init_a, str) and (init_a.lower() == "auto" or init_a.lower() == "init_model"):
@@ -992,33 +991,56 @@ def __init__(
 
         # ### prepare fetch_fn:
         def fetch_fn(idx):
-            X_tensor = tf.py_func(input_data.fetch_X, [idx], dtype)
-            X_tensor.set_shape(
-                idx.get_shape().as_list() + [input_data.num_features]
+            X_tensor = tf.py_func(
+                func=input_data.fetch_X,
+                inp=[idx],
+                Tout=input_data.X.dtype,
+                stateful=False
             )
-
-            design_loc_tensor = tf.py_func(input_data.fetch_design_loc, [idx], dtype)
-            design_loc_tensor.set_shape(
-                idx.get_shape().as_list() + [input_data.num_design_loc_params]
+            X_tensor.set_shape(idx.get_shape().as_list() + [input_data.num_features])
+            X_tensor = tf.cast(X_tensor, dtype=dtype)
+
+            design_loc_tensor = tf.py_func(
+                func=input_data.fetch_design_loc,
+                inp=[idx],
+                Tout=input_data.design_loc.dtype,
+                stateful=False
             )
-            design_scale_tensor = tf.py_func(input_data.fetch_design_scale, [idx], dtype)
-            design_scale_tensor.set_shape(
-                idx.get_shape().as_list() + [input_data.num_design_scale_params]
+            design_loc_tensor.set_shape(idx.get_shape().as_list() + [input_data.num_design_loc_params])
+            design_loc_tensor = tf.cast(design_loc_tensor, dtype=dtype)
+
+            design_scale_tensor = tf.py_func(
+                func=input_data.fetch_design_scale,
+                inp=[idx],
+                Tout=input_data.design_scale.dtype,
+                stateful=False
             )
+            design_scale_tensor.set_shape(idx.get_shape().as_list() + [input_data.num_design_scale_params])
+            design_scale_tensor = tf.cast(design_scale_tensor, dtype=dtype)
 
             if input_data.size_factors is not None:
-                size_factors_tensor = tf.log(tf.py_func(input_data.fetch_size_factors, [idx], dtype))
+                size_factors_tensor = tf.log(tf.py_func(
+                    func=input_data.fetch_size_factors,
+                    inp=[idx],
+                    Tout=input_data.size_factors.dtype,
+                    stateful=False
+                ))
                 size_factors_tensor.set_shape(idx.get_shape())
+                size_factors_tensor = tf.cast(size_factors_tensor, dtype=dtype)
             else:
-                size_factors_tensor = tf.constant(0, shape=(), dtype=X_tensor.dtype)
+                size_factors_tensor = tf.constant(0, shape=(), dtype=dtype)
 
             # return idx, data
             return idx, (X_tensor, design_loc_tensor, design_scale_tensor, size_factors_tensor)
 
         if isinstance(init_a, str):
             init_a = None
+        else:
+            init_a = init_a.astype(dtype)
         if isinstance(init_b, str):
             init_b = None
+        else:
+            init_b = init_b.astype(dtype)
 
         with graph.as_default():
             # create model