Training working fully!!

config.toml

@@ -17,7 +17,7 @@ num_classes = 2 # AD, NL
 [training]
 device = "cuda:0" # "cpu", "cuda", "mps"
 batch_size = 32
-ensemble_size = 50
+ensemble_size = 30
 droprate = 0.05
 learning_rate = 0.0001
-num_epochs = 30
+num_epochs = 25

main.py

@@ -3,6 +3,9 @@ import threading
 import gc
 from typing import Any, Dict
 
+# Must be set before first CUDA context initialization.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
 import toml
 from textual import work
 from textual.app import App, ComposeResult

model/dataset.py

@@ -59,9 +59,11 @@ class ADNIDataset(data.Dataset):  # type: ignore
             mri_data (torch.Tensor): 5D tensor of MRI data with shape (n_samples, channels, width, height, depth).
             xls_data (torch.Tensor): 2D tensor of Excel data with shape (n_samples, features).
         """
-        self.mri_data = mri_data.float().to(device)
-        self.xls_data = xls_data.float().to(device)
-        self.expected_classes = expected_classes.float().to(device)
+        # Keep full datasets on CPU and move only active batches to accelerator.
+        # This avoids loading the entire training corpus into VRAM.
+        self.mri_data = mri_data.float().cpu()
+        self.xls_data = xls_data.float().cpu()
+        self.expected_classes = expected_classes.float().cpu()
         self.filename_ids = filename_ids
 
     def __len__(self) -> int:
@@ -216,8 +218,15 @@ def initalize_dataloaders(
     Returns:
         List[DataLoader[ADNIDataset]]: A list of DataLoaders for the datasets.
     """
+    pin_memory = torch.cuda.is_available()
     return [
-        DataLoader(dataset, batch_size=batch_size, shuffle=True) for dataset in datasets
+        DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            pin_memory=pin_memory,
+        )
+        for dataset in datasets
     ]
 
 

model/dnn_mod.py

@@ -0,0 +1,154 @@
+# Copyright (C) 2024 Intel Labs
+#
+# BSD-3-Clause License
+#
+# Redistribution and use in source and binary forms, with or without modification,
+# are permitted provided that the following conditions are met:
+# 1. Redistributions of source code must retain the above copyright notice,
+#    this list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+# 3. Neither the name of the copyright holder nor the names of its contributors
+#    may be used to endorse or promote products derived from this software
+#    without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS
+# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
+# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
+# OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# Functions related to DNN to BNN model conversion.
+#
+# @authors: Mahesh Subedar
+#
+# ===============================================================================================
+
+from __future__ import absolute_import, division, print_function
+
+import bayesian_torch.layers as bayesian_layers
+from bayesian_torch.utils.util import get_rho
+
+# --------------------------------------------------------------------------------
+# Parameters used to define BNN layyers.
+#    bnn_prior_parameters = {
+#       "prior_mu": 0.0,
+#       "prior_sigma": 1.0,
+#       "posterior_mu_init": 0.0,
+#       "posterior_rho_init": -4.0,
+#       "type": "Reparameterization",  # Flipout or Reparameterization
+# }
+
+
+def bnn_linear_layer(params, d):
+    layer_type = d.__class__.__name__ + params["type"]
+    layer_fn = getattr(bayesian_layers, layer_type)  # Get BNN layer
+    bnn_layer = layer_fn(
+        in_features=d.in_features,
+        out_features=d.out_features,
+        prior_mean=params["prior_mu"],
+        prior_variance=params["prior_sigma"],
+        posterior_mu_init=params["posterior_mu_init"],
+        posterior_rho_init=params["posterior_rho_init"],
+        bias=d.bias is not None,
+    )
+    # if MOPED is enabled initialize mu and sigma
+    if params["moped_enable"]:
+        delta = params["moped_delta"]
+        bnn_layer.mu_weight.data.copy_(d.weight.data)
+        bnn_layer.rho_weight.data.copy_(get_rho(d.weight.data, delta))
+        if bnn_layer.mu_bias is not None:
+            bnn_layer.mu_bias.data.copy_(d.bias.data)
+            bnn_layer.rho_bias.data.copy_(get_rho(d.bias.data, delta))
+    bnn_layer.dnn_to_bnn_flag = True
+    return bnn_layer
+
+
+def bnn_conv_layer(params, d):
+    layer_type = d.__class__.__name__ + params["type"]
+    layer_fn = getattr(bayesian_layers, layer_type)  # Get BNN layer
+    bnn_layer = layer_fn(
+        in_channels=d.in_channels,
+        out_channels=d.out_channels,
+        kernel_size=d.kernel_size,
+        stride=d.stride,
+        padding=d.padding,
+        dilation=d.dilation,
+        groups=d.groups,
+        prior_mean=params["prior_mu"],
+        prior_variance=params["prior_sigma"],
+        posterior_mu_init=params["posterior_mu_init"],
+        posterior_rho_init=params["posterior_rho_init"],
+        bias=d.bias is not None,
+    )
+
+    # if MOPED is enabled,  initialize mu and sigma
+    if params["moped_enable"]:
+        delta = params["moped_delta"]
+        bnn_layer.mu_kernel.data.copy_(d.weight.data)
+        bnn_layer.rho_kernel.data.copy_(get_rho(d.weight.data, delta))
+        if bnn_layer.mu_bias is not None:
+            bnn_layer.mu_bias.data.copy_(d.bias.data)
+            bnn_layer.rho_bias.data.copy_(get_rho(d.bias.data, delta))
+    bnn_layer.dnn_to_bnn_flag = True
+    return bnn_layer
+
+
+def bnn_lstm_layer(params, d):
+    layer_type = d.__class__.__name__ + params["type"]
+    layer_fn = getattr(bayesian_layers, layer_type)  # Get BNN layer
+    bnn_layer = layer_fn(
+        in_features=d.input_size,
+        out_features=d.hidden_size,
+        prior_mean=params["prior_mu"],
+        prior_variance=params["prior_sigma"],
+        posterior_mu_init=params["posterior_mu_init"],
+        posterior_rho_init=params["posterior_rho_init"],
+        bias=d.bias is not None,
+    )
+    # if MOPED is enabled initialize mu and sigma
+    if params["moped_enable"]:
+        print("WARNING: MOPED method is not supported for LSTM layers!!!")
+    bnn_layer.dnn_to_bnn_flag = True
+    return bnn_layer
+
+
+# replaces linear and conv layers
+# bnn_prior_parameters - check the template at the top.
+# MODIFICATION - PREVENT DOUBLE CONVERSION OF DNN TO BNN LAYERS. IF THE LAYER IS ALREADY CONVERTED, SKIP IT.
+def dnn_to_bnn_mod(m, bnn_prior_parameters):
+    for name, value in list(m._modules.items()):
+        # SKIP LAYERS THAT ARE ALREADY CONVERTED TO BNN LAYERS - ENDS IN "Reparameterization" OR "Flipout"
+        if m._modules[name].__class__.__name__.endswith(
+            "Reparameterization"
+        ) or m._modules[name].__class__.__name__.endswith("Flipout"):
+            continue
+        if m._modules[name]._modules:
+            dnn_to_bnn_mod(m._modules[name], bnn_prior_parameters)
+        elif "Conv" in m._modules[name].__class__.__name__:
+            setattr(m, name, bnn_conv_layer(bnn_prior_parameters, m._modules[name]))
+        elif "Linear" in m._modules[name].__class__.__name__:
+            setattr(m, name, bnn_linear_layer(bnn_prior_parameters, m._modules[name]))
+        elif "LSTM" in m._modules[name].__class__.__name__:
+            setattr(m, name, bnn_lstm_layer(bnn_prior_parameters, m._modules[name]))
+        else:
+            pass
+    return
+
+
+def get_kl_loss(m):
+    kl_loss = None
+    for layer in m.modules():
+        if hasattr(layer, "kl_loss"):
+            if kl_loss is None:
+                kl_loss = layer.kl_loss()
+            else:
+                kl_loss += layer.kl_loss()
+    return kl_loss

model/training.py

@@ -21,6 +21,26 @@ type TestMetrics = Tuple[float, float]  # (test_loss, test_acc)
 type KLLossFn = Callable[[nn.Module], torch.Tensor | None]
 
 
+def _model_device(model: nn.Module) -> torch.device:
+    first_param = next(model.parameters(), None)
+    if first_param is None:
+        return torch.device("cpu")
+    return first_param.device
+
+
+def _move_batch_to_model_device(
+    model: nn.Module,
+    mri: torch.Tensor,
+    xls: torch.Tensor,
+    targets: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = _model_device(model)
+    mri = mri.to(device, non_blocking=True)
+    xls = xls.to(device, non_blocking=True)
+    targets = targets.to(device, non_blocking=True)
+    return mri, xls, targets
+
+
 def _batch_correct_and_total(
     outputs: torch.Tensor,
     targets: torch.Tensor,
@@ -84,6 +104,7 @@ def test_model(
     test_progress.update(total=len(test_loader), advance=0)
     for _, (mri, xls, targets, _) in enumerate(test_loader):
         _check_control_events(stop_event=stop_event, pause_event=pause_event)
+        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
         outputs = model((mri, xls))
         loss = criterion(outputs, targets)
         batch_size = mri.size(0)
@@ -138,6 +159,7 @@ def train_epoch(
     train_progress.update(total=len(train_loader), advance=0)
     for _, (mri, xls, targets, _) in enumerate(train_loader):
         _check_control_events(stop_event=stop_event, pause_event=pause_event)
+        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
         optimizer.zero_grad()
         outputs = model((mri, xls))
         loss = criterion(outputs, targets)
@@ -165,6 +187,7 @@ def train_epoch(
     with torch.no_grad():
         for _, (mri, xls, targets, _) in enumerate(val_loader):
             _check_control_events(stop_event=stop_event, pause_event=pause_event)
+            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
             outputs = model((mri, xls))
             loss = criterion(outputs, targets)
             batch_size = mri.size(0)
@@ -292,6 +315,7 @@ def test_model_bayesian(
     with torch.no_grad():
         for _, (mri, xls, targets, _) in enumerate(test_loader):
             _check_control_events(stop_event=stop_event, pause_event=pause_event)
+            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
             outputs = model((mri, xls))
             data_loss = cast(torch.Tensor, criterion(outputs, targets))
             batch_size = mri.size(0)
@@ -343,6 +367,7 @@ def train_epoch_bayesian(
     train_progress.update(total=len(train_loader), advance=0)
     for _, (mri, xls, targets, _) in enumerate(train_loader):
         _check_control_events(stop_event=stop_event, pause_event=pause_event)
+        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
         optimizer.zero_grad()
         outputs = model((mri, xls))
         data_loss = cast(torch.Tensor, criterion(outputs, targets))
@@ -377,6 +402,7 @@ def train_epoch_bayesian(
     with torch.no_grad():
         for _, (mri, xls, targets, _) in enumerate(val_loader):
             _check_control_events(stop_event=stop_event, pause_event=pause_event)
+            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
             outputs = model((mri, xls))
             data_loss = cast(torch.Tensor, criterion(outputs, targets))
             batch_size = mri.size(0)

tasks/train_bayesian.py

@@ -1,17 +1,17 @@
+import gc
 import pathlib as pl
 import time
-import gc
 from threading import Event
 from typing import Any, Dict
 
 import torch
-from bayesian_torch.models.dnn_to_bnn import dnn_to_bnn, get_kl_loss  # type: ignore[import-untyped]
 from torch import nn, optim
 from torch.utils.data import DataLoader
 
 import model.dataset as ds
 import model.training as tn
 from model.cnn import CNN3D
+from model.dnn_mod import dnn_to_bnn_mod, get_kl_loss
 from util.progress import ProgressTracker
 from util.ui_logger import PipelineLogger
 
@@ -78,17 +78,15 @@ def train_bayesian_task(
         log.info(
             f"Training model {model_num + 1}/{config['training']['ensemble_size']}..."
         )
-        model = (
-            CNN3D(
-                image_channels=config["data"]["image_channels"],
-                clin_data_channels=config["data"]["clin_data_channels"],
-                num_classes=config["data"]["num_classes"],
-                droprate=config["training"]["droprate"],
-            )
-            .float()
-            .to(config["training"]["device"])
-        )
-        dnn_to_bnn(model, bnn_prior_parameters)
+        model = CNN3D(
+            image_channels=config["data"]["image_channels"],
+            clin_data_channels=config["data"]["clin_data_channels"],
+            num_classes=config["data"]["num_classes"],
+            droprate=config["training"]["droprate"],
+        ).float()
+
+        dnn_to_bnn_mod(model, bnn_prior_parameters)
+        model.to(config["training"]["device"])
 
         optimizer = optim.Adam(
             model.parameters(), lr=config["training"]["learning_rate"]
@@ -123,8 +121,7 @@ def train_bayesian_task(
         )
 
         log.info(
-            f"Run {model_num + 1}/{config['training']['ensemble_size']} - "
-            f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.4f}"
+            f"Run {model_num + 1}/{config['training']['ensemble_size']} - Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.4f}"
         )
 
         model_save_path = bayesian_models_path / f"model_{model_num + 1}.pt"