added eager executing tf2 training with nb noise

Author: picciama

batchglm/train/tf2/__init__.py

@@ -0,0 +1,3 @@
+from . import glm_nb as nb
+from . import glm_norm as norm
+from . import glm_beta as beta

batchglm/train/tf2/base/__init__.py

@@ -0,0 +1,3 @@
+from .estimator import TFEstimator
+from .model import ProcessModelBase, ModelBase, LossBase
+from .optim import OptimizerBase

batchglm/train/tf2/base/estimator.py

@@ -0,0 +1,32 @@
+from .external import pkg_constants, TrainingStrategies
+from .model import ModelBase, LossBase
+
+import numpy as np
+import tensorflow as tf
+
+
+class TFEstimator:
+    model: ModelBase
+    loss: LossBase
+
+    def __init__(self, input_data, dtype):
+
+        self._input_data = input_data
+        self.dtype = dtype
+
+    def _train(
+            self,
+            batched_model: bool,
+            batch_size: int,
+            optimizer_object: tf.keras.optimizers.Optimizer,
+            optimizer_enum: TrainingStrategies,
+            convergence_criteria: str,
+            stopping_criteria: int,
+            autograd: bool,
+            featurewise: bool,
+            benchmark: bool
+    ):
+        pass
+
+    def fetch_fn(self, idx):
+        pass

batchglm/train/tf2/base/external.py

@@ -0,0 +1,5 @@
+#from batchglm.models.base import _Estimator_Base
+#from batchglm.xarray_sparse import SparseXArrayDataArray, SparseXArrayDataSet
+from batchglm.train.tf2.base_glm.training_strategies import TrainingStrategies
+#import batchglm.utils.stats as stat_utils
+from batchglm import pkg_constants

batchglm/train/tf2/base/model.py

@@ -0,0 +1,57 @@
+import abc
+import logging
+import tensorflow as tf
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class ModelBase(tf.keras.Model, metaclass=abc.ABCMeta):
+
+    def __init__(self):
+        super(ModelBase, self).__init__()
+
+    @abc.abstractmethod
+    def call(self, inputs, training=False, mask=None):
+        pass
+
+
+class LossBase(tf.keras.losses.Loss, metaclass=abc.ABCMeta):
+
+    def __init__(self):
+        super(LossBase, self).__init__()
+
+    @abc.abstractmethod
+    def call(self, y_true, y_pred):
+        pass
+
+
+class ProcessModelBase:
+
+    @abc.abstractmethod
+    def param_bounds(self, dtype):
+        pass
+
+    def tf_clip_param(
+            self,
+            param,
+            name
+    ):
+        bounds_min, bounds_max = self.param_bounds(param.dtype)
+        return tf.clip_by_value(
+            param,
+            bounds_min[name],
+            bounds_max[name]
+        )
+
+    def np_clip_param(
+            self,
+            param,
+            name
+    ):
+        bounds_min, bounds_max = self.param_bounds(param.dtype)
+        return np.clip(
+            param,
+            bounds_min[name],
+            bounds_max[name]
+        )

batchglm/train/tf2/base/optim.py

@@ -0,0 +1,52 @@
+import abc
+import logging
+import tensorflow as tf
+
+logger = logging.getLogger("batchglm")
+
+
+class OptimizerBase(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):
+
+    def __init__(self, name):
+        super(OptimizerBase, self).__init__(name=name)
+
+    @abc.abstractmethod
+    def _resource_apply_dense(self, grad, handle):
+        pass
+
+    @abc.abstractmethod
+    def _resource_apply_sparse(self, grad, handle, apply_state):
+        pass
+
+    @abc.abstractmethod
+    def _create_slots(self):
+        pass
+
+    """
+    @property
+    @abc.abstractmethod
+    def vars(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def gradients(self):
+        return None
+
+    @property
+    @abc.abstractmethod
+    def hessians(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def fims(self):
+        pass
+
+    @abc.abstractmethod
+    def step(self, learning_rate):
+        pass
+    """
+    @abc.abstractmethod
+    def get_config(self):
+        pass

batchglm/train/tf2/base_glm/README.md

@@ -0,0 +1,2 @@
+# Classes with GLM specific code.
+All noise models that are in the GLM category inherit all of these classes.

batchglm/train/tf2/base_glm/__init__.py

@@ -0,0 +1,10 @@
+from .processModel import ProcessModelGLM
+from .model import GLM, LossGLM
+
+from .estimator import Estimator
+from .vars import ModelVarsGLM
+from .layers import LinearLocGLM, LinearScaleGLM, LinkerLocGLM, LinkerScaleGLM
+from .layers import LikelihoodGLM, UnpackParamsGLM
+from .layers_gradients import JacobianGLM, HessianGLM, FIMGLM
+from .optim import NR, IRLS
+from .training_strategies import TrainingStrategies