Tensorflow 2.12 installation

requirements/tensorflow_2_11_requirements.txt

@@ -21,4 +21,5 @@ tf_siren
 cupy==12.0.0
 scikit-cuda
 yellowbrick
-seaborn
+seaborn
+pathos

requirements/tensorflow_2_12_requirements.txt

@@ -0,0 +1,25 @@
+tensorflow==2.12.*
+tensorflow-addons==0.20.0
+tensorflow-probability==0.20.1
+mrcfile==1.4.3
+h5py
+pyyaml==5.1
+bibtexparser<=1.2.0
+psutil<=5.8.0
+tkcolorpicker<=2.1.3
+tifffile==2021.3.5
+matplotlib==3.6.2
+scipy==1.9.3
+pandas
+biopython==1.76
+scikit-learn==1.1.3
+numpy==1.23.4
+protobuf==3.20.3
+xmipp_metadata@git+https://github.com/DavidHerreros/xmipp_metadata.git@devel
+prody
+tf_siren
+cupy-cuda11x
+scikit-cuda
+yellowbrick
+seaborn
+pathos

requirements/tensorflow_2_3_requirements.txt

@@ -20,4 +20,5 @@ tf_siren
 cupy-cuda101
 scikit-cuda
 yellowbrick
-seaborn
+seaborn
+pathos

requirements/tensorflow_2_4_requirements.txt

@@ -20,4 +20,5 @@ tf_siren
 cupy==12.0.0
 scikit-cuda
 yellowbrick
-seaborn
+seaborn
+pathos

setup.py

@@ -5,7 +5,6 @@
 https://github.com/pypa/sampleproject
 """
 
-
 import os
 import sys
 import subprocess
@@ -58,8 +57,12 @@ def condaInstallationCommands(self):
 
         # If at least one GPU is series 3000 and above, change installation requirements
         for gpu_model in gpu_models:
-            if re.findall(r"30[0-9]+", gpu_model) or version.parse("450.80.02") <= version.parse(driver):
-                cuda_version = "11"
+            if re.findall(r"40[0-9]+", gpu_model) or re.findall(r"30[0-9]+", gpu_model) and \
+                    version.parse("520.56.06") <= version.parse(driver):
+                cuda_version = "11.8"
+                break
+            elif re.findall(r"30[0-9]+", gpu_model) or version.parse("450.80.02") <= version.parse(driver):
+                cuda_version = "11.2"
                 break
 
         self.print_flush("Cuda version to be installed: " + cuda_version)
@@ -71,7 +74,12 @@ def condaInstallationCommands(self):
         condabin_path_command = r"which conda | sed 's: ::g'"
 
         # Command: Get installation of new conda env with Cuda, Cudnn, and Tensorflow dependencies
-        if cuda_version == "11":
+        if cuda_version == "11.8":
+            req_file = os.path.join("requirements", "tensorflow_2_12_requirements.txt")
+            command = "if ! { conda env list | grep 'flexutils-tensorflow'; } >/dev/null 2>&1; then " \
+                      "conda create -y -n flexutils-tensorflow " \
+                      "-c conda-forge python=3.9 cudatoolkit=11.8 cudatoolkit-dev pyyaml -y; fi"
+        elif cuda_version == "11.2":
             req_file = os.path.join("requirements", "tensorflow_2_11_requirements.txt")
             command = "if ! { conda env list | grep 'flexutils-tensorflow'; } >/dev/null 2>&1; then " \
                       "conda create -y -n flexutils-tensorflow " \
@@ -82,7 +90,7 @@ def condaInstallationCommands(self):
                       "conda create -y -n flexutils-tensorflow -c conda-forge python=3.8 cudatoolkit=10.1 cudnn=7" \
                       "cudatoolkit-dev pyyaml -y; fi"
 
-        return req_file, conda_path_command, condabin_path_command, command
+        return req_file, conda_path_command, condabin_path_command, command, cuda_version
 
     def runCondaInstallation(self):
         # Check conda is in PATH
@@ -100,22 +108,54 @@ def runCondaInstallation(self):
         self.installMissingPackages()
         self.print_flush("...done")
 
-        req_file, _, condabin_path_command, install_conda_command = self.condaInstallationCommands()
+        req_file, _, condabin_path_command, install_conda_command, cuda_version = self.condaInstallationCommands()
 
         self.print_flush("Installing Tensorflow conda env...")
         subprocess.check_call(install_conda_command, shell=True)
         self.print_flush("...done")
 
         self.print_flush("Getting env pip...")
         path = subprocess.check_output(condabin_path_command, shell=True).decode("utf-8").replace('\n', '').replace("*", "")
-        install_toolkit_command = 'eval "$(%s shell.bash hook)" && conda activate flexutils-tensorflow && ' \
-                                  'pip install -r %s && pip install -e toolkit' % (path, req_file)
+        if cuda_version == "11.8":
+            install_toolkit_command = 'eval "$(%s shell.bash hook)" && conda activate flexutils-tensorflow && ' \
+                                      'conda install -c nvidia cuda-nvcc=11.3.58 && ' \
+                                      'pip install nvidia-cudnn-cu11==8.6.0.163 && ' \
+                                      'pip install -r %s && pip install -e toolkit' % (path, req_file)
+
+        else:
+            install_toolkit_command = 'eval "$(%s shell.bash hook)" && conda activate flexutils-tensorflow && ' \
+                                      'pip install -r %s && pip install -e toolkit' % (path, req_file)
         self.print_flush("...done")
 
         self.print_flush("Installing Flexutils-Tensorflow toolkit in conda env...")
         subprocess.check_call(install_toolkit_command, shell=True)
         self.print_flush("...done")
 
+        self.print_flush("Set environment variables in conda env...")
+        if cuda_version == "11.8":
+            commands = ['eval "$(%s shell.bash hook) "' % path,
+                        'conda activate flexutils-tensorflow ',
+                        'mkdir -p $CONDA_PREFIX/etc/conda/activate.d ',
+                        'echo \'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))\''
+                        ' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
+                        'echo \'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib\' '
+                        '>> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
+                        'echo \'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/lib/\n\' >> '
+                        '$CONDA_PREFIX/etc/conda/activate.d/env_vars.sh ',
+                        'mkdir -p $CONDA_PREFIX/lib/nvvm/libdevice ',
+                        'cp $CONDA_PREFIX/lib/libdevice.10.bc $CONDA_PREFIX/lib/nvvm/libdevice/'
+                        ]
+        else:
+            commands = ['eval "$(%s shell.bash hook) "' % path,
+                        'conda activate flexutils-tensorflow ',
+                        'mkdir -p $CONDA_PREFIX/etc/conda/activate.d ',
+                        'echo \'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/\' '
+                        '>> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh'
+                        ]
+        commands = "&&".join(commands)
+        subprocess.check_call(commands, shell=True)
+        self.print_flush("...done")
+
 
 here = os.path.abspath(os.path.dirname(__file__))
 
@@ -132,8 +172,8 @@ def runCondaInstallation(self):
       keywords='scipion continuous-heterogeneity imageprocessing xmipp',
       packages=find_packages(),
       package_data={  # Optional
-         'requirements': ["*"],
-         'toolkit': ["*"]
+          'requirements': ["*"],
+          'toolkit': ["*"]
       },
       cmdclass={'install': Installation}
       )

toolkit/tensorflow_toolkit/scripts/predict_deep_nma.py

@@ -52,7 +52,8 @@ def predict(md_file, weigths_file, n_modes, refinePose, architecture, ctfType, p
 
     # Load model
     autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Predict step
     print("------------------ Predicting NMA coefficients... ------------------")

toolkit/tensorflow_toolkit/scripts/predict_deep_pose.py

@@ -49,7 +49,8 @@ def predict(md_file, weigths_file, refinePose, architecture, ctfType, pad=2, sr=
 
     # Load model
     autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Get poses
     # alignment = []

toolkit/tensorflow_toolkit/scripts/predict_flex_consensus.py

@@ -89,7 +89,8 @@ def predict(outPath, dataPath, weigths_file, latDim):
 
     # Load model
     autoencoder = AutoEncoder(generator)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Predict step
     data = spaces[0]

toolkit/tensorflow_toolkit/scripts/predict_het_siren.py

@@ -55,7 +55,8 @@ def predict(md_file, weigths_file, refinePose, architecture, ctfType, pad=2, sr=
     # Load model
     autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType, refPose=refinePose,
                               het_dim=hetDim)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Get poses
     print("------------------ Predicting particles... ------------------")
@@ -83,7 +84,7 @@ def predict(md_file, weigths_file, refinePose, architecture, ctfType, pad=2, sr=
 
     # Save map
     for idx, decoded_map in enumerate(decoded_maps):
-        decoded_path = Path(Path(md_file).parent, 'decoded_map_class_%d.mrc' % (idx + 1))
+        decoded_path = Path(Path(md_file).parent, 'decoded_map_class_%02d.mrc' % (idx + 1))
         ImageHandler().write(decoded_map, decoded_path, overwrite=True)
 
 

toolkit/tensorflow_toolkit/scripts/predict_homo_siren.py

@@ -53,7 +53,8 @@ def predict(md_file, weigths_file, refinePose, architecture, ctfType, pad=2, sr=
 
     # Load model
     autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType, refPose=refinePose)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Get poses
     print("------------------ Predicting particles... ------------------")

toolkit/tensorflow_toolkit/scripts/predict_map_het_siren.py

@@ -56,12 +56,13 @@ def predict(weigths_file, het_file, out_path, allCoords=False, filter=False, **k
 
     # Load model
     autoencoder = AutoEncoder(generator, het_dim=x_het.shape[1], **kwargs)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Decode maps
     decoded_maps = autoencoder.eval_volume_het(x_het, allCoords=allCoords, filter=filter)
     for idx, decoded_map in enumerate(decoded_maps):
-        decoded_path = Path(out_path, 'decoded_map_class_%d.mrc' % (idx + 1))
+        decoded_path = Path(out_path, 'decoded_map_class_%02d.mrc' % (idx + 1))
         ImageHandler().write(decoded_map, decoded_path, overwrite=True)
 
 

toolkit/tensorflow_toolkit/scripts/predict_particles_het_siren.py

@@ -56,7 +56,8 @@ def predict(md_file, weigths_file, refinePose, architecture, ctfType,
     # Load model
     autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType, refPose=refinePose,
                               het_dim=hetDim)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Get poses
     print("------------------ Predicting alignment and het info... ------------------")

toolkit/tensorflow_toolkit/scripts/predict_zernike3deep.py

@@ -49,7 +49,8 @@ def predict(md_file, weigths_file, L1, L2, refinePose, architecture, ctfType, pa
 
     # Load model
     autoencoder = AutoEncoder(generator, architecture=architecture, CTF=ctfType)
-    autoencoder.load_weights(weigths_file).expect_partial()
+    autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
+    autoencoder.load_weights(weigths_file)
 
     # Get Zernike3DSpace
     # zernike_space = []

toolkit/tensorflow_toolkit/scripts/train_deep_nma.py

@@ -65,7 +65,7 @@ def train(outPath, md_file, n_modes, batch_size, shuffle, splitTrain, epochs, co
         raise error
 
     # Save model
-    autoencoder.save_weights(os.path.join(outPath, "deep_nma_model"))
+    autoencoder.save_weights(os.path.join(outPath, "deep_nma_model.h5"))
 
 
 if __name__ == '__main__':

toolkit/tensorflow_toolkit/scripts/train_deep_pose.py

@@ -59,6 +59,7 @@ def train(outPath, md_file, batch_size, shuffle, step, splitTrain, epochs, cost,
 
         # Fine tune a previous model
         if weigths_file:
+            autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
             autoencoder.load_weights(weigths_file)
 
         # if tf.__version__ == '2.11.0':
@@ -78,7 +79,7 @@ def train(outPath, md_file, batch_size, shuffle, step, splitTrain, epochs, cost,
         raise error
 
     # Save model
-    autoencoder.save_weights(os.path.join(outPath, "deep_pose_model"))
+    autoencoder.save_weights(os.path.join(outPath, "deep_pose_model.h5"))
 
 
 if __name__ == '__main__':

toolkit/tensorflow_toolkit/scripts/train_flex_consensus.py

@@ -92,7 +92,7 @@ def train(outPath, dataPath, latDim, batch_size, shuffle, splitTrain, epochs):
         raise error
 
     # Save model
-    autoencoder.save_weights(os.path.join(outPath, "network", "flex_consensus_model"))
+    autoencoder.save_weights(os.path.join(outPath, "network", "flex_consensus_model.h5"))
 
     # Get templates for future matching
     best_mean = [None, None]

toolkit/tensorflow_toolkit/scripts/train_het_siren.py

@@ -57,6 +57,7 @@ def train(outPath, md_file, batch_size, shuffle, step, splitTrain, epochs, cost,
 
         # Fine tune a previous model
         if weigths_file:
+            autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
             autoencoder.load_weights(weigths_file)
 
         optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
@@ -71,7 +72,7 @@ def train(outPath, md_file, batch_size, shuffle, step, splitTrain, epochs, cost,
         raise error
 
     # Save model
-    autoencoder.save_weights(os.path.join(outPath, "het_siren_model"))
+    autoencoder.save_weights(os.path.join(outPath, "het_siren_model.h5"))
 
 
 if __name__ == '__main__':

toolkit/tensorflow_toolkit/scripts/train_homo_siren.py

@@ -57,6 +57,7 @@ def train(outPath, md_file, batch_size, shuffle, step, splitTrain, epochs, cost,
 
         # Fine tune a previous model
         if weigths_file:
+            autoencoder.build(input_shape=(None, generator.xsize, generator.xsize, 1))
             autoencoder.load_weights(weigths_file)
 
         if superConv:
@@ -74,7 +75,7 @@ def train(outPath, md_file, batch_size, shuffle, step, splitTrain, epochs, cost,
         raise error
 
     # Save model
-    autoencoder.save_weights(os.path.join(outPath, "homo_siren_model"))
+    autoencoder.save_weights(os.path.join(outPath, "homo_siren_model.h5"))
 
 
 if __name__ == '__main__':

toolkit/tensorflow_toolkit/scripts/train_zernike3deep.py

@@ -67,7 +67,7 @@ def train(outPath, md_file, L1, L2, batch_size, shuffle, step, splitTrain, epoch
         raise error
 
     # Save model
-    autoencoder.save_weights(os.path.join(outPath, "zernike3deep_model"))
+    autoencoder.save_weights(os.path.join(outPath, "zernike3deep_model.h5"))
 
 
 if __name__ == '__main__':