Merge pull request #78 from AllenInstitute/fix/gpu_memory

move to predict_on_batch

deepinterpolation/__init__.py

@@ -1,2 +1,2 @@
 # Version of deepinterpolation package. Should match the version in setup.py
-__version__ = "0.1.4"
+__version__ = "0.1.5"

deepinterpolation/cli/schemas.py

@@ -468,6 +468,26 @@ class TrainingSchema(argschema.schemas.DefaultSchema):
             checkpoints.",
     )
 
+    use_multiprocessing = argschema.fields.Bool(
+        required=False,
+        default=True,
+        description="whether to use a multiprocessing pool to fetch batch \
+            samples. Setting this to true will increase data generation speed \
+            if the generator is limited by read speed. This will also \
+            increase RAM memory usage. Set to False if your hardware \
+            encounter RAM memory error during training.",
+    )
+
+    nb_workers = argschema.fields.Int(
+        required=False,
+        default=16,
+        description="Nb of workers running on the CPU to fetch \
+            batch samples. This parameter is only relevant if \
+            use_multiprocessing is set to True. Larger number of nb_workers \
+            will increase memory usage. Increase this number until your \
+            training becomes limited either by RAM or CPU usage.",
+    )
+
 
 class FineTuningSchema(argschema.schemas.DefaultSchema):
     name = argschema.fields.String(
@@ -574,6 +594,26 @@ class FineTuningSchema(argschema.schemas.DefaultSchema):
             checkpoints.",
     )
 
+    use_multiprocessing = argschema.fields.Bool(
+        required=False,
+        default=True,
+        description="whether to use a multiprocessing pool to fetch batch \
+            samples. Setting this to true will increase data generation speed \
+            if the generator is limited by read speed. This will also \
+            increase RAM memory usage. Set to False if your hardware \
+            encounter RAM memory error during training.",
+    )
+
+    nb_workers = argschema.fields.Int(
+        required=False,
+        default=16,
+        description="Nb of workers running on the CPU to fetch \
+            batch samples. This parameter is only relevant if \
+            use_multiprocessing is set to True. Larger number of nb_workers \
+            will increase memory usage. Increase this number until your \
+            training becomes limited either by RAM or CPU usage.",
+    )
+
 
 class NetworkSchema(argschema.schemas.DefaultSchema):
     name = argschema.fields.String(

deepinterpolation/inferrence_collection.py

@@ -243,7 +243,7 @@ def run(self):
             for index_dataset in np.arange(0, self.nb_datasets, 1):
                 local_data = self.generator_obj.__getitem__(index_dataset)
 
-                predictions_data = self.model.predict(local_data[0])
+                predictions_data = self.model.predict_on_batch(local_data[0])
 
                 local_mean, local_std = \
                     self.generator_obj.__get_norm_parameters__(index_dataset)

requirements.txt

@@ -1,6 +1,5 @@
-tensorflow==2.4.4
+tensorflow==2.7
 nibabel
-h5py==2.10.0
 matplotlib
 numpy
 python-dateutil

setup.py

@@ -14,7 +14,7 @@
 
 setup(
     name="deepinterpolation",
-    version="0.1.4",
+    version="0.1.5",
     description="Implemenent DeepInterpolation to denoise data by removing \
 independent noise",
     long_description_content_type='text/x-rst',