Add interface for model initialization

R/Estimator.R

@@ -470,12 +470,14 @@ createEstimator <- function(modelParameters,
   path <- system.file("python", package = "DeepPatientLevelPrediction")
 
   if (modelParameters$modelType == "Finetuner") {
-    estimatorSettings$finetune <- TRUE
     plpModel <- PatientLevelPrediction::loadPlpModel(modelParameters$modelPath)
     estimatorSettings$finetuneModelPath <-
       normalizePath(file.path(plpModel$model, "DeepEstimatorModel.pt"))
     modelParameters$modelType <-
       plpModel$modelDesign$modelSettings$modelType
+    initStrategy <- reticulate::import_from_path("InitStrategy", path = path)$FinetuneInitStrategy()
+  } else {
+    initStrategy <- reticulate::import_from_path("InitStrategy", path = path)$DefaultInitStrategy()
   }
 
   model <-
@@ -490,7 +492,8 @@ createEstimator <- function(modelParameters,
   estimator <- estimator(
     model = model,
     model_parameters = modelParameters,
-    estimator_settings = estimatorSettings
+    estimator_settings = estimatorSettings,
+    init_strategy = initStrategy
   )
   return(estimator)
 }

inst/python/Estimator.py

@@ -6,31 +6,23 @@
 from tqdm import tqdm
 
 from gpu_memory_cleanup import memory_cleanup
-
+from InitStrategy import InitStrategy, DefaultInitStrategy
 
 class Estimator:
     """
     A class that wraps around pytorch models.
     """
 
-    def __init__(self, model, model_parameters, estimator_settings):
+    def __init__(self, model, model_parameters, estimator_settings, init_strategy: InitStrategy = DefaultInitStrategy()):
         self.seed = estimator_settings["seed"]
         if callable(estimator_settings["device"]):
             self.device = estimator_settings["device"]()
         else:
             self.device = estimator_settings["device"]
         torch.manual_seed(seed=self.seed)
-        if "finetune" in estimator_settings.keys() and estimator_settings["finetune"]:
-            path = estimator_settings["finetune_model_path"]
-            fitted_estimator = torch.load(path, map_location="cpu")
-            fitted_parameters = fitted_estimator["model_parameters"]
-            self.model = model(**fitted_parameters)
-            self.model.load_state_dict(fitted_estimator["model_state_dict"])
-            for param in self.model.parameters():
-                param.requires_grad = False
-            self.model.reset_head()
-        else:
-            self.model = model(**model_parameters)
+
+        self.model = init_strategy.initialize(model, model_parameters, estimator_settings)
+
         self.model_parameters = model_parameters
         self.estimator_settings = estimator_settings
 

inst/python/InitStrategy.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+import torch
+
+class InitStrategy(ABC):
+    @abstractmethod
+    def initialize(self, model, model_parameters, estimator_settings):
+        pass
+
+class DefaultInitStrategy(InitStrategy):
+    def initialize(self, model, model_parameters, estimator_settings):
+        return model(**model_parameters)
+
+class FinetuneInitStrategy(InitStrategy):
+    def initialize(self, model, model_parameters, estimator_settings):
+        path = estimator_settings["finetune_model_path"]
+        fitted_estimator = torch.load(path, map_location="cpu")
+        fitted_parameters = fitted_estimator["model_parameters"]
+        model_instance = model(**fitted_parameters)
+        model_instance.load_state_dict(fitted_estimator["model_state_dict"])
+        for param in model_instance.parameters():
+            param.requires_grad = False
+        model_instance.reset_head()
+        return model_instance
+

tests/testthat/test-Finetuner.R

@@ -44,7 +44,6 @@ test_that("Finetuner fitEstimator works", {
 
   fineTunedModel <- torch$load(file.path(fineTunerResults$model,
                                          "DeepEstimatorModel.pt"))
-  expect_true(fineTunedModel$estimator_settings$finetune)
   expect_equal(fineTunedModel$estimator_settings$finetune_model_path, 
                normalizePath(file.path(fitEstimatorPath, "plpModel", "model",
                          "DeepEstimatorModel.pt")))