mlcommons · Geeks-Sid · Mar 21, 2022 · Nov 4, 2021 · Nov 4, 2021 · Nov 4, 2021
@@ -23,6 +23,7 @@ jobs:
     - name: Install dependencies and package
       run: |
         python -m pip install --upgrade pip
+        python -m pip install openvino-dev
         $CONDA/bin/conda install -c conda-forge libvips -y
         pip3 install torch==1.8.2+cpu torchvision==0.9.2+cpu torchaudio==0.8.2 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
         pip install -e .

@@ -9,6 +9,7 @@ RUN python3.7 -m pip install --upgrade pip
 RUN python3.7 -m pip install torch==1.10.0+cpu torchvision==0.11.0+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html
 COPY . /GaNDLF
 WORKDIR /GaNDLF
+RUN python3.7 -m pip install --upgrade pip && python3.7 -m pip install openvino-dev
 RUN python3.7 -m pip install -e .
 # Entrypoint forces all commands given via "docker run" to go through python, CMD forces the default entrypoint script argument to be gandlf_run
 # If a user calls "docker run gandlf:[tag] gandlf_anonymize", it will resolve to running "python gandlf_anonymize" instead.

@@ -9,6 +9,7 @@ LABEL version=1.0
 RUN python3 -m pip install --upgrade pip
 COPY . /GaNDLF
 WORKDIR /GaNDLF
+RUN python3 -m pip install --upgrade pip && python3 -m pip install openvino-dev
 RUN python3 -m pip install -e .
 # Entrypoint forces all commands given via "docker run" to go through python, CMD forces the default entrypoint script argument to be gandlf_run
 # If a user calls "docker run gandlf:[tag] gandlf_anonymize", it will resolve to running "python gandlf_anonymize" instead.

@@ -12,6 +12,7 @@ RUN python3.7 -m pip install --upgrade pip
 RUN python3.7 -m pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
 COPY . /GaNDLF
 WORKDIR /GaNDLF
+RUN python3.7 -m pip install --upgrade pip && python3.7 -m pip install openvino-dev
 RUN python3.7 -m pip install -e .
 # Entrypoint forces all commands given via "docker run" to go through python, CMD forces the default entrypoint script argument to be gandlf_run
 # If a user calls "docker run gandlf:[tag] gandlf_anonymize", it will resolve to running "python gandlf_anonymize" instead.

@@ -9,6 +9,7 @@ LABEL version=1.0
 RUN python3 -m pip install --upgrade pip
 COPY . /GaNDLF
 WORKDIR /GaNDLF
+RUN python3 -m pip install --upgrade pip && python3 -m pip install openvino-dev
 RUN python3 -m pip install -e .
 # Entrypoint forces all commands given via "docker run" to go through python, CMD forces the default entrypoint script argument to be gandlf_run
 # If a user calls "docker run gandlf:[tag] gandlf_anonymize", it will resolve to running "python gandlf_anonymize" instead.

@@ -24,7 +24,7 @@ def validate_network(
 
     Parameters
     ----------
-    model : torch.model
+    model : if parameters["model"]["type"] == torch, this is a torch.model, otherwise this is OV exec_net
         The model to process the input image with, it should support appropriate dimensions.
     valid_dataloader : torch.DataLoader
         The dataloader for the validation epoch
@@ -76,11 +76,13 @@ def validate_network(
     pathlib.Path(current_output_dir).mkdir(parents=True, exist_ok=True)
 
     # Set the model to valid
-    model.eval()
+    if params["model"]["type"] == "torch":
+        model.eval()
+
     # # putting stuff in individual arrays for correlation analysis
     # all_targets = []
     # all_predics = []
-    if params["medcam_enabled"]:
+    if params["medcam_enabled"] and params["model"]["type"] == "torch":
         model.enable_medcam()
         params["medcam_enabled"] = True
 
@@ -158,12 +160,23 @@ def validate_network(
                 ## special case for 2D
                 if image.shape[-1] == 1:
                     image = torch.squeeze(image, -1)
-                pred_output += model(image)
+                if params["model"]["type"] == "torch":
+                    pred_output += model(image)
+                elif params["model"]["type"] == "openvino":
+                    pred_output += torch.from_numpy(
+                        model.infer(
+                            inputs={params["model"]["IO"][0]: image.cpu().numpy()}
+                        )[params["model"]["IO"][1]]
+                    )
+                else:
+                    raise Exception(
+                        "Model type not supported. Please only use 'torch' or 'openvino'."
+                    )
+
             pred_output = pred_output.cpu() / params["q_samples_per_volume"]
             pred_output /= params["scaling_factor"]
             # all_predics.append(pred_output.double())
             # all_targets.append(valuesToPredict.double())
-            print(f"pred_output.shape: {pred_output.shape}")
 
             if is_inference and is_classification:
                 logits_list.append(pred_output)
@@ -242,7 +255,10 @@ def validate_network(
                         flush=True,
                     )
 
-                result = step(model, image, label, params)
+                if is_inference:
+                    result = step(model, image, label, params, train=False)
+                else:
+                    result = step(model, image, label, params, train=True)
 
                 # get the current attention map and add it to its aggregator
                 if params["medcam_enabled"]:
@@ -326,7 +342,7 @@ def validate_network(
                     )
 
             # get the final attention map and save it
-            if params["medcam_enabled"]:
+            if params["medcam_enabled"] and params["model"]["type"] == "torch":
                 attention_map = attention_map_aggregator.get_output_tensor()
                 for i, n in enumerate(attention_map):
                     model.save_attention_map(
@@ -383,7 +399,7 @@ def validate_network(
                     to_print,
                 )
 
-    if params["medcam_enabled"]:
+    if params["medcam_enabled"] and params["model"]["type"] == "torch":
         model.disable_medcam()
         params["medcam_enabled"] = False
 

@@ -15,8 +15,12 @@
 import tiffslide as openslide
 
 from GANDLF.data.ImagesFromDataFrame import ImagesFromDataFrame
-from GANDLF.utils import populate_channel_keys_in_params, send_model_to_device
 from GANDLF.models import global_models_dict
+from GANDLF.utils import (
+    populate_channel_keys_in_params,
+    send_model_to_device,
+    load_ov_model,
+)
 from GANDLF.data.inference_dataloader_histopath import InferTumorSegDataset
 
 
@@ -31,6 +35,7 @@ def inference_loop(inferenceDataFromPickle, device, parameters, outputDir):
         outputDir (str): The output directory.
     """
     # Defining our model here according to parameters mentioned in the configuration file
+    print("Current model type : ", parameters["model"]["type"])
     print("Number of dims     : ", parameters["model"]["dimension"])
     if "num_channels" in parameters["model"]:
         print("Number of channels : ", parameters["model"]["num_channels"])
@@ -47,17 +52,46 @@ def inference_loop(inferenceDataFromPickle, device, parameters, outputDir):
     )
     inference_loader = DataLoader(inferenceDataForTorch, batch_size=1)
 
-    # Loading the weights into the model
-    main_dict = outputDir
-    if os.path.isdir(outputDir):
-        file_to_check = os.path.join(
-            outputDir, str(parameters["model"]["architecture"]) + "_best.pth.tar"
+    if parameters["model"]["type"] == "torch":
+        # Loading the weights into the model
+        main_dict = outputDir
+        if os.path.isdir(outputDir):
+            file_to_check = os.path.join(
+                outputDir, str(parameters["model"]["architecture"]) + "_best.pth.tar"
+            )
+            if not os.path.isfile(file_to_check):
+                raise ValueError(
+                    "The specified model was not found: {0}.".format(file_to_check)
+                )
+
+        main_dict = torch.load(file_to_check, map_location=torch.device(device))
+        model.load_state_dict(main_dict["model_state_dict"])
+    elif parameters["model"]["type"].lower() == "openvino":
+        # Loading the executable OpenVINO model
+        main_dict = outputDir
+        if os.path.isdir(outputDir):
+            xml_to_check = os.path.join(
+                outputDir, str(parameters["model"]["architecture"]) + "_best.xml"
+            )
+            bin_to_check = os.path.join(
+                outputDir, str(parameters["model"]["architecture"]) + "_best.bin"
+            )
+            if not os.path.isfile(xml_to_check):
+                raise ValueError(
+                    "The specified model IR was not found: {0}.".format(xml_to_check)
+                )
+            if not os.path.isfile(bin_to_check):
+                raise ValueError(
+                    "The model specified model weights was not found: {0}.".format(
+                        bin_to_check
+                    )
+                )
+            model, input_blob, output_blob = load_ov_model(xml_to_check, device.upper())
+            parameters["model"]["IO"] = [input_blob, output_blob]
+    else:
+        raise ValueError(
+            "The model type is not recognized: ", parameters["model"]["type"]
         )
-        if not os.path.isfile(file_to_check):
-            raise ValueError("The model specified model was not found:", file_to_check)
-
-    main_dict = torch.load(file_to_check, map_location=torch.device(device))
-    model.load_state_dict(main_dict["model_state_dict"])
 
     if not (os.environ.get("HOSTNAME") is None):
         print("\nHostname     :" + str(os.environ.get("HOSTNAME")), flush=True)
@@ -67,9 +101,10 @@ def inference_loop(inferenceDataFromPickle, device, parameters, outputDir):
     parameters["save_output"] = True
 
     print("Data Samples: ", len(inference_loader.dataset), flush=True)
-    model, parameters["model"]["amp"], parameters["device"] = send_model_to_device(
-        model, parameters["model"]["amp"], device, optimizer=None
-    )
+    if parameters["model"]["type"] == "torch":
+        model, parameters["model"]["amp"], parameters["device"] = send_model_to_device(
+            model, parameters["model"]["amp"], device, optimizer=None
+        )
 
     print("Using device:", parameters["device"], flush=True)
 
@@ -129,12 +164,23 @@ def inference_loop(inferenceDataFromPickle, device, parameters, outputDir):
             )
             for image_patches, (x_coords, y_coords) in tqdm(dataloader):
                 x_coords, y_coords = y_coords.numpy(), x_coords.numpy()
-                if parameters["model"]["amp"]:
-                    with autocast():
+                if parameters["model"]["type"] == "torch":
+                    if parameters["model"]["amp"]:
+                        with autocast():
+                            output = model(
+                                image_patches.float().to(parameters["device"])
+                            )
+                    else:
                         output = model(image_patches.float().to(parameters["device"]))
+                    output = output.detach().cpu().numpy()
                 else:
-                    output = model(image_patches.float().to(parameters["device"]))
-                output = output.detach().cpu().numpy()
+                    output = model.infer(
+                        inputs={
+                            parameters["model"]["IO"][0]: image_patches.float()
+                            .cpu()
+                            .numpy()
+                        }
+                    )[parameters["model"]["IO"][1]]
                 for i in range(int(output.shape[0])):
                     count_map[
                         x_coords[i] : x_coords[i] + patch_size[0],

@@ -3,7 +3,7 @@
 from .loss_and_metric import get_loss_and_metrics
 
 
-def step(model, image, label, params):
+def step(model, image, label, params, train=True):
     """
     Function that steps the model for a single batch
 
@@ -60,11 +60,19 @@ def step(model, image, label, params):
             if len(label.shape) > 1:
                 label = torch.squeeze(label, -1)
 
-    if params["model"]["amp"]:
-        with torch.cuda.amp.autocast():
-            output = model(image)
+    if train == False and params["model"]["type"].lower() == "openvino":
+        output = torch.from_numpy(
+            model.infer(inputs={params["model"]["IO"][0]: image.cpu().numpy()})[
+                params["model"]["IO"][1]
+            ]
+        )
+        output = output.to(params["device"])
     else:
-        output = model(image)
+        if params["model"]["amp"]:
+            with torch.cuda.amp.autocast():
+                output = model(image)
+        else:
+            output = model(image)
 
     if "medcam_enabled" in params and params["medcam_enabled"]:
         output, attention_map = output

@@ -358,17 +358,16 @@ def training_loop(
 
     # if previous model file is present, load it up
     if os.path.exists(best_model_path):
-        print("Previous model found. Loading it up.")
         try:
             main_dict = load_model(best_model_path)
             version_check(params["version"], version_to_check=main_dict["version"])
             model.load_state_dict(main_dict["model_state_dict"])
             start_epoch = main_dict["epoch"]
             optimizer.load_state_dict(main_dict["optimizer_state_dict"])
             best_loss = main_dict["loss"]
-            print("Previous model loaded successfully.")
-        except IOError:
-            raise IOError("Previous model could not be loaded, error: ")
+            print("Previous model successfully loaded.")
+        except RuntimeWarning:
+            RuntimeWarning("Previous model could not be loaded, initializing model")
 
     print("Using device:", device, flush=True)
 
@@ -460,15 +459,21 @@ def training_loop(
             best_loss = epoch_valid_loss
             best_train_idx = epoch
             patience = 0
+
+            model.eval()
             save_model(
                 {
                     "epoch": best_train_idx,
                     "model_state_dict": model.state_dict(),
                     "optimizer_state_dict": optimizer.state_dict(),
                     "loss": best_loss,
                 },
+                model,
+                params,
                 best_model_path,
+                onnx_export=False,
             )
+            model.train()
             first_model_saved = True
 
         print("Current Best epoch: ", best_train_idx)
@@ -491,6 +496,42 @@ def training_loop(
         flush=True,
     )
 
+    # once the training is done, optimize the best model
+    if os.path.exists(best_model_path):
+
+        onnx_export = True
+        if params["model"]["architecture"] in ["sdnet", "brain_age"]:
+            onnx_export = False
+        elif (
+            "onnx_export" in params["model"] and params["model"]["onnx_export"] == False
+        ):
+            onnx_export = False
+
+        if onnx_export:
+            print("Optimizing best model.")
+
+            try:
+ 
10000
               main_dict = load_model(best_model_path)
+                version_check(params["version"], version_to_check=main_dict["version"])
+                model.load_state_dict(main_dict["model_state_dict"])
+                best_epoch = main_dict["epoch"]
+                optimizer.load_state_dict(main_dict["optimizer_state_dict"])
+                best_loss = main_dict["loss"]
+                save_model(
+                    {
+                        "epoch": best_epoch,
+                        "model_state_dict": model.state_dict(),
+                        "optimizer_state_dict": optimizer.state_dict(),
+                        "loss": best_loss,
+                    },
+                    model,
+                    params,
+                    best_model_path,
+                    onnx_export,
+                )
+            except Exception as e:
+                print("Best model could not be loaded, error: ", e)
+
 
 if __name__ == "__main__":
 

@@ -38,6 +38,10 @@ def InferenceManager(dataframe, outputDir, parameters, device):
     class_list = None
     is_classification = parameters["problem_type"] == "classification"
 
+    # initialize model type for processing: if not defined, default to torch
+    if not ("type" in parameters["model"]):
+        parameters["model"]["type"] = "torch"
+
     for fold_dir in fold_dirs:
         parameters["current_fold_dir"] = fold_dir
         inference_loop(

@@ -29,6 +29,6 @@ def forward(self, x):
         if isinstance(B, int):
             return F.avg_pool3d(x, (W, H, D)).view(B, C)
         else:
-            return F.avg_pool2d(x, (W.item(), H.item(), D.item())).view(
+            return F.avg_pool3d(x, (W.item(), H.item(), D.item())).view(
                 B.item(), C.item()
             )