Tensorflow Version update to 2.19 (#597)

* .h5 to .keras and many more updates!

* python 3.11 and 3.12

---------

Co-authored-by: rxrathod <rrathod@broadinstitute.org>

Author: lucidtronix

.github/workflows/python-package.yml

@@ -21,7 +21,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9"]
+        python-version: [3.11, 3.12]
 
     steps:
     - uses: actions/checkout@v2
@@ -35,7 +35,6 @@ jobs:
         python -m pip install --upgrade pip
         # Install the ml4h Python package.
         pip install .
-        pip install -r docker/vm_boot_images/config/tensorflow-requirements.txt
-    - name: Test with pytest
+    - name: Test with pytest and pytest-xdist
       run: |
-        pytest tests -m "not slow"
+        pytest tests -m "not slow" -n auto

docker/vm_boot_images/Dockerfile

@@ -20,6 +20,11 @@ WORKDIR /app
 # to minimize full recompilation where possible.
 
 # Basic setup
+#RUN rm /etc/apt/sources.list.d/cuda.list
+##RUN rm /etc/apt/sources.list.d/nvidia-ml.list
+#RUN apt-key del 7fa2af80
+#RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+#RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
 RUN ./ubuntu.sh
 
 # Point any MLflow tracking hooks at the main MLflow instance on Cloud Run
@@ -32,7 +37,10 @@ ENV MLFLOW_TRACKING_URI='https://mlflow-783282864357.us-central1.run.app'
 # RUN pip3 uninstall -y fastai
 # RUN ./fastai.sh
 
-RUN apt-get install python3-tk libgl1-mesa-glx libxt-dev -y
+RUN apt-get update
+RUN apt-get upgrade -y
+RUN apt-get install python3 python3-pip python3-tk libgl1-mesa-glx libxt-dev -y
+RUN apt-get install -y wget unzip curl python3-pydot graphviz git ffmpeg
 
 # Requirements for the tensorflow project
 RUN pip3 install --upgrade pip

docker/vm_boot_images/build.sh

@@ -19,11 +19,11 @@ CPU_ONLY="false"
 PUSH_TO_GCR="false"
 PUSH_TO_LATEST="false"
 
-BASE_IMAGE_GPU="tensorflow/tensorflow:2.9.1-gpu"
-BASE_IMAGE_CPU="tensorflow/tensorflow:2.9.1"
+BASE_IMAGE_GPU="tensorflow/tensorflow:2.19.0-gpu"
+BASE_IMAGE_CPU="tensorflow/tensorflow:2.19.0"
 
-LATEST_TAG_GPU="tf2.9-latest-gpu"
-LATEST_TAG_CPU="tf2.9-latest-cpu"
+LATEST_TAG_GPU="tf2.19-latest-gpu"
+LATEST_TAG_CPU="tf2.19-latest-cpu"
 
 SCRIPT_NAME=$( echo $0 | sed 's#.*/##g' )
 

docker/vm_boot_images/config/tensorflow-requirements.txt

@@ -1,23 +1,21 @@
 pydot
 nibabel==4.0.2
 pydicom==1.2.2
-hyperopt==0.1.2
 seaborn
 scikit-image
 peakutils
 biosppy
-vtk==9.2.6
-imageio==2.6.1
+imageio
 ipywidgets>=7.5.1
 bokeh
 pillow
 notebook
 pytest
+pytest-xdist
 pysam
-tensorflow==2.9.1
-tensorflow-addons==0.17.1
+tensorflow==2.19.0
 tensorflow_hub
-tensorflow_probability==0.17.0
+tensorflow_probability
 tensorflow-text
 tf-models-official
 keras-tuner
@@ -35,7 +33,7 @@ plotnine
 vega
 ipycanvas>=0.7.0
 ipyannotations>=0.2.1
-torch==1.12.1
+torch==2.2.2
 opencv-python
 blosc
 boto3
@@ -47,4 +45,3 @@ voxelmorph
 pystrum
 av
 lmdb
-mlflow

docker/vm_boot_images/config/ubuntu.sh

@@ -2,5 +2,13 @@
 
 # Other necessities
 apt-get update
+
 echo "ttf-mscorefonts-installer msttcorefonts/accepted-mscorefonts-eula select true" | debconf-set-selections
-apt-get install -y wget unzip curl python3-pydot python3-pydot-ng graphviz ttf-mscorefonts-installer git pip ffmpeg hdf5-tools
+
+apt-get install -y wget unzip curl python3-pydot python3-pydot-ng graphviz ttf-mscorefonts-installer git pip ffmpeg
+
+wget https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb
+dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb
+cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-local-8138232B-keyring.gpg /usr/share/keyrings/
+apt-get update
+apt-get -y install cudnn

ml4h/arguments.py

@@ -28,6 +28,9 @@
 from ml4h.TensorMap import TensorMap, TimeSeriesOrder
 from ml4h.defines import IMPUTATION_RANDOM, IMPUTATION_MEAN
 from ml4h.tensormap.mgb.dynamic import make_mgb_dynamic_tensor_maps
+from ml4h.tensormap.tensor_map_maker import generate_categorical_tensor_map_from_file, \
+    generate_latent_tensor_map_from_file
+
 from ml4h.models.legacy_models import parent_sort, check_no_bottleneck
 from ml4h.tensormap.tensor_map_maker import make_test_tensor_maps, generate_random_pixel_as_text_tensor_maps
 from ml4h.models.legacy_models import NORMALIZATION_CLASSES, CONV_REGULARIZATION_CLASSES, DENSE_REGULARIZATION_CLASSES
@@ -292,7 +295,7 @@ def parse_args():
         help='Maximum number of models for the hyper-parameter optimizer to evaluate before returning.',
     )
     parser.add_argument('--balance_csvs', default=[], nargs='*', help='Balances batches with representation from sample IDs in this list of CSVs')
-    parser.add_argument('--optimizer', default='radam', type=str, help='Optimizer for model training')
+    parser.add_argument('--optimizer', default='adam', type=str, help='Optimizer for model training')
     parser.add_argument('--learning_rate_schedule', default=None, type=str, choices=['triangular', 'triangular2', 'cosine_decay'], help='Adjusts learning rate during training.')
     parser.add_argument('--anneal_rate', default=0., type=float, help='Annealing rate in epochs of loss terms during training')
     parser.add_argument('--anneal_shift', default=0., type=float, help='Annealing offset in epochs of loss terms during training')
@@ -430,6 +433,10 @@ def parse_args():
     # TensorMap prefix for convenience
     parser.add_argument('--tensormap_prefix', default="ml4h.tensormap", type=str, help="Module prefix path for TensorMaps. Defaults to \"ml4h.tensormap\"")
 
+    #Parent Sort enable or disable
+    parser.add_argument('--parent_sort', default=True, type=lambda x: x.lower() == 'true', help='disable or enable parent_sort on output tmaps')
+    #Dictionary outputs
+    parser.add_argument('--named_outputs', default=False, type=lambda x: x.lower() == 'true', help='pass output tmaps as dictionaries if true else pass as list')
     args = parser.parse_args()
     _process_args(args)
     return args
@@ -538,7 +545,7 @@ def _process_args(args):
 
     if args.latent_input_file is not None:
         args.tensor_maps_in.append(
-            generate_latent_tensor_map_from_file(args.latent_input_file, args.input_tensors.pop(0)),
+            generate_latent_tensor_map_from_file(args.latent_input_file, args.input_tensors.pop(0))
         )
     args.tensor_maps_in.extend([tensormap_lookup(it, args.tensormap_prefix) for it in args.input_tensors])
 
@@ -564,11 +571,13 @@ def _process_args(args):
                     args.output_tensors.pop(0),
                 ),
             )
+
     if len(args.latent_output_files) > 0:
         for lof in args.latent_output_files:
             args.tensor_maps_out.append(
                 generate_latent_tensor_map_from_file(lof, args.output_tensors.pop(0)),
             )
+
     args.tensor_maps_out.extend([tensormap_lookup(ot, args.tensormap_prefix) for ot in args.output_tensors])
     args.tensor_maps_out = parent_sort(args.tensor_maps_out)
     args.tensor_maps_protected = [tensormap_lookup(it, args.tensormap_prefix) for it in args.protected_tensors]

ml4h/defines.py

@@ -28,7 +28,7 @@ def __str__(self):
 IMAGE_EXT = '.png'
 PDF_EXT = '.pdf'
 TENSOR_EXT = '.hd5'
-MODEL_EXT = '.h5'
+MODEL_EXT = '.keras'
 XML_EXT = '.xml'
 
 STOP_CHAR = '!'

ml4h/metrics.py

@@ -8,15 +8,14 @@
 
 from sklearn.metrics import roc_curve, auc, average_precision_score
 
-
 from tensorflow.keras.losses import binary_crossentropy, categorical_crossentropy, sparse_categorical_crossentropy
-from tensorflow.keras.losses import logcosh, cosine_similarity, mean_squared_error, mean_absolute_error, mean_absolute_percentage_error
+from tensorflow.keras.losses import LogCosh, CosineSimilarity, MSE, MAE, MAPE, Dice
+from keras.saving import register_keras_serializable
 
-from neurite.tf.losses import Dice
 
 STRING_METRICS = [
     'categorical_crossentropy','binary_crossentropy','mean_absolute_error','mae',
-    'mean_squared_error', 'mse', 'cosine_similarity', 'logcosh', 'sparse_categorical_crossentropy',
+    'mean_squared_error', 'mse', 'cosine_similarity', 'log_cosh', 'sparse_categorical_crossentropy',
 ]
 
 
@@ -48,6 +47,7 @@ def weighted_crossentropy(weights, name='anonymous'):
     string_fxn += '\treturn loss\n'
     exec(string_fxn, globals(), locals())
     loss_fxn = eval(name + fxn_postfix, globals(), locals())
+    loss_fxn = register_keras_serializable()(loss_fxn)
     return loss_fxn
 
 
@@ -109,39 +109,39 @@ def paired_angle_between_batches(tensors):
 
 def ignore_zeros_l2(y_true, y_pred):
     mask = K.cast(K.not_equal(y_true, 0), K.floatx())
-    return mean_squared_error(y_true * mask, y_pred * mask)
+    return MSE(y_true * mask, y_pred * mask)
 
 
 def ignore_zeros_logcosh(y_true, y_pred):
     mask = K.cast(K.not_equal(y_true, 0), K.floatx())
-    return logcosh(y_true * mask, y_pred * mask)
+    return LogCosh(y_true * mask, y_pred * mask)
 
 
 def sentinel_logcosh_loss(sentinel: float):
     def ignore_sentinel_logcosh(y_true, y_pred):
         mask = K.cast(K.not_equal(y_true, sentinel), K.floatx())
-        return logcosh(y_true * mask, y_pred * mask)
+        return LogCosh(y_true * mask, y_pred * mask)
     return ignore_sentinel_logcosh
 
 
 def y_true_times_mse(y_true, y_pred):
-    return K.maximum(y_true, 1.0)*mean_squared_error(y_true, y_pred)
+    return K.maximum(y_true, 1.0)*MSE(y_true, y_pred)
 
 
 def mse_10x(y_true, y_pred):
-    return 10.0*mean_squared_error(y_true, y_pred)
+    return 10.0*MSE(y_true, y_pred)
 
 
 def y_true_squared_times_mse(y_true, y_pred):
-    return K.maximum(1.0+y_true, 1.0)*K.maximum(1.0+y_true, 1.0)*mean_squared_error(y_true, y_pred)
+    return K.maximum(1.0+y_true, 1.0)*K.maximum(1.0+y_true, 1.0)*MSE(y_true, y_pred)
 
 
 def y_true_cubed_times_mse(y_true, y_pred):
-    return K.maximum(y_true, 1.0)*K.maximum(y_true, 1.0)*K.maximum(y_true, 1.0)*mean_squared_error(y_true, y_pred)
+    return K.maximum(y_true, 1.0)*K.maximum(y_true, 1.0)*K.maximum(y_true, 1.0)*MSE(y_true, y_pred)
 
 
 def y_true_squared_times_logcosh(y_true, y_pred):
-    return K.maximum(1.0+y_true, 1.0)*K.maximum(1.0+y_true, 1.0)*logcosh(y_true, y_pred)
+    return K.maximum(1.0+y_true, 1.0)*K.maximum(1.0+y_true, 1.0)*LogCosh(y_true, y_pred)
 
 
 def two_batch_euclidean(tensors):
@@ -265,6 +265,7 @@ def loss(y_true, y_pred):
     return loss
 
 def dice(y_true, y_pred):
+    return Dice()(y_true, y_pred)
     return Dice(laplace_smoothing=1e-05).mean_loss(y_true, y_pred)
 
 def per_class_dice(labels):
@@ -273,12 +274,13 @@ def per_class_dice(labels):
         label_idx = labels[label_key]
         fxn_name = label_key.replace('-', '_').replace(' ', '_')
         string_fxn = 'def ' + fxn_name + '_dice(y_true, y_pred):\n'
-        string_fxn += '\tdice = Dice(laplace_smoothing=1e-05).dice(y_true, y_pred)\n'
-        string_fxn += '\tdice = K.mean(dice, axis=0)['+str(label_idx)+']\n'
+        string_fxn += '\tdice = tf.keras.losses.Dice()(y_true, y_pred)\n'
+        #string_fxn += '\tdice = K.mean(dice, axis=0)['+str(label_idx)+']\n'
         string_fxn += '\treturn dice'
 
         exec(string_fxn)
         dice_fxn = eval(fxn_name + '_dice')
+        dice_fxn = register_keras_serializable()(dice_fxn)
         dice_fxns.append(dice_fxn)
 
     return dice_fxns
@@ -299,6 +301,7 @@ def per_class_recall(labels):
 
         exec(string_fxn)
         recall_fxn = eval(fxn_name + '_recall')
+        recall_fxn = register_keras_serializable()(recall_fxn)
         recall_fxns.append(recall_fxn)
 
     return recall_fxns
@@ -317,6 +320,7 @@ def per_class_precision(labels):
 
         exec(string_fxn)
         precision_fxn = eval(fxn_name + '_precision')
+        precision_fxn = register_keras_serializable()(precision_fxn)
         precision_fxns.append(precision_fxn)
 
     return precision_fxns
@@ -335,6 +339,7 @@ def per_class_recall_3d(labels):
 
         exec(string_fxn)
         recall_fxn = eval(fxn_prefix + '_recall')
+        recall_fxn = register_keras_serializable()(recall_fxn)
         recall_fxns.append(recall_fxn)
 
     return recall_fxns
@@ -353,6 +358,7 @@ def per_class_precision_3d(labels):
 
         exec(string_fxn)
         precision_fxn = eval(fxn_prefix + '_precision')
+        precision_fxn = register_keras_serializable()(precision_fxn)
         precision_fxns.append(precision_fxn)
 
     return precision_fxns
@@ -371,6 +377,7 @@ def per_class_recall_4d(labels):
 
         exec(string_fxn)
         recall_fxn = eval(fxn_prefix + '_recall')
+        recall_fxn = register_keras_serializable()(recall_fxn)
         recall_fxns.append(recall_fxn)
 
     return recall_fxns
@@ -389,6 +396,8 @@ def per_class_precision_4d(labels):
 
         exec(string_fxn)
         precision_fxn = eval(fxn_prefix + '_precision')
+        precision_fxn = register_keras_serializable()(precision_fxn)
+
         precision_fxns.append(precision_fxn)
 
     return precision_fxns
@@ -407,6 +416,7 @@ def per_class_recall_5d(labels):
 
         exec(string_fxn)
         recall_fxn = eval(fxn_prefix + '_recall')
+        recall_fxn = register_keras_serializable()(recall_fxn)
         recall_fxns.append(recall_fxn)
 
     return recall_fxns
@@ -425,6 +435,7 @@ def per_class_precision_5d(labels):
 
         exec(string_fxn)
         precision_fxn = eval(fxn_prefix + '_precision')
+        precision_fxn = register_keras_serializable()(precision_fxn)
         precision_fxns.append(precision_fxn)
 
     return precision_fxns
@@ -449,15 +460,15 @@ def get_metric_dict(output_tensor_maps):
         elif tm.loss == 'binary_crossentropy':
             losses.append(binary_crossentropy)
         elif tm.loss == 'mean_absolute_error' or tm.loss == 'mae':
-            losses.append(mean_absolute_error)
+            losses.append(MSE)
         elif tm.loss == 'mean_squared_error' or tm.loss == 'mse':
-            losses.append(mean_squared_error)
+            losses.append(MSE)
         elif tm.loss == 'cosine_similarity':
-            losses.append(cosine_similarity)
-        elif tm.loss == 'logcosh':
-            losses.append(logcosh)
+            losses.append(CosineSimilarity)
+        elif tm.loss == 'log_cosh':
+            losses.append(LogCosh)
         elif tm.loss == 'mape':
-            losses.append(mean_absolute_percentage_error)
+            losses.append(MAPE)
         elif hasattr(tm.loss,  '__name__'):
             metrics[tm.loss.__name__] = tm.loss
             losses.append(tm.loss)
@@ -857,4 +868,12 @@ def result(self):
     def reset_state(self):
         # Reset the metric state variables
         self.total_ssim.assign(0.0)
-        self.count.assign(0.0)
+        self.count.assign(0.0)
+
+        
+def _register_all(module_globals):
+    for name, obj in module_globals.items():
+        if callable(obj) and not name.startswith("_"):
+            module_globals[name] = register_keras_serializable()(obj)
+
+_register_all(globals())