Allows to clear backends model

drivendataorg · Feb 1, 2018 · 34494f6 · 34494f6
1 parent 32e002d
commit 34494f6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/prediction/src/algorithms/classify/src/classification_model.py b/prediction/src/algorithms/classify/src/classification_model.py
@@ -117,3 +117,10 @@ def predict(self, candidates, model_path=None):
                                'p_concerning': float}}
         """
         pass
+
+    @abc.abstractmethod
+    def clear(self):
+        """
+        Clear all signs of backend models.
+        """
+        pass
diff --git a/prediction/src/algorithms/classify/src/lr3dcnn/model.py b/prediction/src/algorithms/classify/src/lr3dcnn/model.py
@@ -2,9 +2,10 @@
 import numpy as np
 import os
 from collections import defaultdict
+from keras import backend as K
+from src.algorithms.classify.src.classification_model import ClassificationModel
 from src.preprocess import preprocess_ct, load_ct, crop_patches, generators
 
-from src.algorithms.classify.src.classification_model import ClassificationModel
 from .architecture import net
 
 
@@ -307,3 +308,9 @@ def predict(self, candidates, model_path=None):
                 iter += 1
 
         return candidates
+
+    def clear(self):
+        del self.model
+        self.pull_ct.clear()
+        self.pull_patches.clear()
+        K.clear_session()
diff --git a/prediction/src/tests/test_classification_3dlrcnn.py b/prediction/src/tests/test_classification_3dlrcnn.py
@@ -1,6 +1,7 @@
-from ..algorithms.classify.src.lr3dcnn.model import Model
 import os
 
+from ..algorithms.classify.src.lr3dcnn.model import Model
+
 
 def test_classify_init_load(models_dir_path):
     model = Model(init_model=True, pull_size=10, batch_size=32, data_format=None)
@@ -12,6 +13,7 @@ def test_classify_init_load(models_dir_path):
         data_format=None
     )
     assert model is not None
+    model.clear()
 
 
 def test_classify_predict(dicom_paths, nodule_locations, models_dir_path):
@@ -24,6 +26,7 @@ def test_classify_predict(dicom_paths, nodule_locations, models_dir_path):
     candidates = [{'file_path': dicom_paths[0], 'centroids': [nl for nl in nodule_locations]}]
     predicted = model.predict(candidates)
     assert 0 <= predicted[0]['centroids'][0]['p_concerning'] <= 1
+    model.clear()
 
 
 def test_classify_real_nodule_full_dicom(dicom_paths, models_dir_path):
@@ -37,3 +40,4 @@ def test_classify_real_nodule_full_dicom(dicom_paths, models_dir_path):
     predicted = model.predict(candidates)
     print(predicted)
     assert .3 <= predicted[0]['centroids'][0]['p_concerning'] <= 1
+    model.clear()