Upload 3523 files

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/MNIST_Classification/data/FashionMNIST/raw/t10k-images-idx3-ubyte filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/MNIST_Classification/data/FashionMNIST/raw/train-images-idx3-ubyte filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/MNIST_Classification/data/MNIST/raw/t10k-images-idx3-ubyte filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/MNIST_Classification/data/MNIST/raw/train-images-idx3-ubyte filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/UCI_Benchmarks/data/uci/concrete/Concrete_Data.xls filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/UCI_Benchmarks/data/uci/power-plant/Folds5x2_pp.ods filter=lfs diff=lfs merge=lfs -text
+SPC-UQ/UCI_Benchmarks/data/uci/power-plant/Folds5x2_pp.xlsx filter=lfs diff=lfs merge=lfs -text

SPC-UQ/.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

SPC-UQ/.idea/UQ_baseline.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="py37" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

SPC-UQ/.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,24 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="11">
+            <item index="0" class="java.lang.String" itemvalue="tqdm" />
+            <item index="1" class="java.lang.String" itemvalue="scipy" />
+            <item index="2" class="java.lang.String" itemvalue="tabulate" />
+            <item index="3" class="java.lang.String" itemvalue="scikit_learn" />
+            <item index="4" class="java.lang.String" itemvalue="matplotlib" />
+            <item index="5" class="java.lang.String" itemvalue="gpytorch" />
+            <item index="6" class="java.lang.String" itemvalue="torch" />
+            <item index="7" class="java.lang.String" itemvalue="setuptools" />
+            <item index="8" class="java.lang.String" itemvalue="numpy" />
+            <item index="9" class="java.lang.String" itemvalue="torchvision" />
+            <item index="10" class="java.lang.String" itemvalue="Pillow" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

SPC-UQ/.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

SPC-UQ/.idea/misc.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Black">
+    <option name="sdkName" value="py37" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="py37" project-jdk-type="Python SDK" />
+</project>

SPC-UQ/.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/UQ_baseline.iml" filepath="$PROJECT_DIR$/.idea/UQ_baseline.iml" />
+    </modules>
+  </component>
+</project>

SPC-UQ/.idea/other.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="PySciProjectComponent">
+    <option name="PY_INTERACTIVE_PLOTS_SUGGESTED" value="true" />
+  </component>
+</project>

SPC-UQ/.idea/workspace.xml

@@ -0,0 +1,247 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="AutoImportSettings">
+    <option name="autoReloadType" value="SELECTIVE" />
+  </component>
+  <component name="ChangeListManager">
+    <list default="true" id="5a477d09-bea8-4806-81a6-0e58ccd074ce" name="Changes" comment="" />
+    <option name="SHOW_DIALOG" value="false" />
+    <option name="HIGHLIGHT_CONFLICTS" value="true" />
+    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
+    <option name="LAST_RESOLUTION" value="IGNORE" />
+  </component>
+  <component name="FileTemplateManagerImpl">
+    <option name="RECENT_TEMPLATES">
+      <list>
+        <option value="Python Script" />
+      </list>
+    </option>
+  </component>
+  <component name="ProjectColorInfo">{
+  &quot;associatedIndex&quot;: 2
+}</component>
+  <component name="ProjectId" id="30Vf7yOdfmRNkiAM7wBXYZUSVri" />
+  <component name="ProjectViewState">
+    <option name="hideEmptyMiddlePackages" value="true" />
+    <option name="showLibraryContents" value="true" />
+  </component>
+  <component name="PropertiesComponent">{
+  &quot;keyToString&quot;: {
+    &quot;Python.evidential.executor&quot;: &quot;Run&quot;,
+    &quot;Python.re.executor&quot;: &quot;Run&quot;,
+    &quot;Python.rename.executor&quot;: &quot;Run&quot;,
+    &quot;Python.run_cls_tests.executor&quot;: &quot;Run&quot;,
+    &quot;Python.run_cubic_tests.executor&quot;: &quot;Run&quot;,
+    &quot;Python.run_toy_tests.executor&quot;: &quot;Run&quot;,
+    &quot;Python.run_uci_dataset_tests (1).executor&quot;: &quot;Run&quot;,
+    &quot;Python.run_uci_dataset_tests (2).executor&quot;: &quot;Run&quot;,
+    &quot;Python.run_uci_dataset_tests.executor&quot;: &quot;Run&quot;,
+    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
+    &quot;last_opened_file_path&quot;: &quot;E:/Experiment/SPC-UQ/Depth_regression/trainers&quot;,
+    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
+    &quot;node.js.detected.package.tslint&quot;: &quot;true&quot;,
+    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
+    &quot;node.js.selected.package.tslint&quot;: &quot;(autodetect)&quot;,
+    &quot;nodejs_package_manager_path&quot;: &quot;npm&quot;,
+    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
+  }
+}</component>
+  <component name="RecentsManager">
+    <key name="CopyFile.RECENT_KEYS">
+      <recent name="E:\Experiment\SPC-UQ\Depth_regression\trainers" />
+    </key>
+    <key name="MoveFile.RECENT_KEYS">
+      <recent name="E:\Experiment\SPC-UQ\Depth_regression" />
+    </key>
+  </component>
+  <component name="RunManager" selected="Python.run_cls_tests">
+    <configuration name="run_cls_tests" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="UQ_baseline" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/MNIST_Classification" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/MNIST_Classification/run_cls_tests.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="run_cubic_tests" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="UQ_baseline" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/Cubic_Regression" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/Cubic_Regression/run_cubic_tests.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="run_toy_tests" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="UQ_baseline" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/Toy_regression" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/Toy_regression/run_toy_tests.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="run_uci_dataset_tests (1)" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="UQ_baseline" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/UCI_Benchmarks" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/UCI_Benchmarks/run_uci_dataset_tests.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="run_uci_dataset_tests (2)" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="UQ_baseline" />
+      <option name="ENV_FILES" value="" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/../UCI_Benchmarks" />
+      <option name="IS_MODULE_SDK" value="false" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/../UCI_Benchmarks/run_uci_dataset_tests.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <recent_temporary>
+      <list>
+        <item itemvalue="Python.run_cls_tests" />
+        <item itemvalue="Python.run_cubic_tests" />
+        <item itemvalue="Python.run_uci_dataset_tests (2)" />
+        <item itemvalue="Python.run_uci_dataset_tests (1)" />
+        <item itemvalue="Python.run_toy_tests" />
+      </list>
+    </recent_temporary>
+  </component>
+  <component name="SharedIndexes">
+    <attachedChunks>
+      <set>
+        <option value="bundled-js-predefined-1d06a55b98c1-91d5c284f522-JavaScript-PY-241.15989.155" />
+        <option value="bundled-python-sdk-babbdf50b680-7c6932dee5e4-com.jetbrains.pycharm.pro.sharedIndexes.bundled-PY-241.15989.155" />
+      </set>
+    </attachedChunks>
+  </component>
+  <component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="application-level" UseSingleDictionary="true" transferred="true" />
+  <component name="TaskManager">
+    <task active="true" id="Default" summary="Default task">
+      <changelist id="5a477d09-bea8-4806-81a6-0e58ccd074ce" name="Changes" comment="" />
+      <created>1753717498209</created>
+      <option name="number" value="Default" />
+      <option name="presentableId" value="Default" />
+      <updated>1753717498209</updated>
+      <workItem from="1753717499301" duration="61000" />
+      <workItem from="1753717566735" duration="10589000" />
+      <workItem from="1753789230404" duration="17167000" />
+      <workItem from="1753911784054" duration="686000" />
+      <workItem from="1753975576642" duration="9813000" />
+      <workItem from="1754180731050" duration="28000" />
+      <workItem from="1754221817091" duration="34326000" />
+      <workItem from="1754311631259" duration="1625000" />
+      <workItem from="1754430709706" duration="5414000" />
+      <workItem from="1754511717348" duration="1628000" />
+      <workItem from="1754839988337" duration="2408000" />
+      <workItem from="1754919749833" duration="13485000" />
+      <workItem from="1755004077990" duration="3927000" />
+      <workItem from="1756907577926" duration="3062000" />
+      <workItem from="1756995689716" duration="6564000" />
+      <workItem from="1757086350108" duration="1819000" />
+      <workItem from="1757163460686" duration="1011000" />
+    </task>
+    <servers />
+  </component>
+  <component name="TypeScriptGeneratedFilesManager">
+    <option name="version" value="3" />
+  </component>
+  <component name="XDebuggerManager">
+    <breakpoint-manager>
+      <default-breakpoints>
+        <breakpoint type="python-exception">
+          <properties notifyOnTerminate="true" exception="BaseException">
+            <option name="notifyOnTerminate" value="true" />
+          </properties>
+        </breakpoint>
+      </default-breakpoints>
+    </breakpoint-manager>
+  </component>
+  <component name="com.intellij.coverage.CoverageDataManagerImpl">
+    <SUITE FILE_PATH="coverage/UQ_baseline$evidential.coverage" NAME="evidential Coverage Results" MODIFIED="1753831104933" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/Toy_classification/trainers" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$run_toy_tests.coverage" NAME="run_toy_tests Coverage Results" MODIFIED="1754511868967" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/Toy_regression" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$rename.coverage" NAME="rename Coverage Results" MODIFIED="1754235285059" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/UCI_regression/trainers" />
+    <SUITE FILE_PATH="coverage/UQ_baseline$run_toy_tests.coverage" NAME="run_toy_tests Coverage Results" MODIFIED="1753976501175" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/Toy_regression" />
+    <SUITE FILE_PATH="coverage/UQ_baseline$run_cls_tests.coverage" NAME="run_cls_tests Coverage Results" MODIFIED="1753837116769" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/Toy_classification" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$run_uci_dataset_tests__1_.coverage" NAME="run_uci_dataset_tests (1) Coverage Results" MODIFIED="1756909439135" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/UCI_Benchmarks" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$run_uci_dataset_tests__2_.coverage" NAME="run_uci_dataset_tests (2) Coverage Results" MODIFIED="1756909699750" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/../UCI_Benchmarks" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$run_uci_dataset_tests.coverage" NAME="run_uci_dataset_tests Coverage Results" MODIFIED="1754253016762" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/UCI_regression" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$re.coverage" NAME="re Coverage Results" MODIFIED="1754263188586" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/Depth_regression/save" />
+    <SUITE FILE_PATH="coverage/UQ_baseline$run_uci_dataset_tests.coverage" NAME="run_uci_dataset_tests Coverage Results" MODIFIED="1753995527450" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/UCI_regression" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$run_cubic_tests.coverage" NAME="run_cubic_tests Coverage Results" MODIFIED="1756996456492" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/Cubic_Regression" />
+    <SUITE FILE_PATH="coverage/SPC_UQ$run_cls_tests.coverage" NAME="run_cls_tests Coverage Results" MODIFIED="1757002010234" SOURCE_PROVIDER="com.intellij.coverage.DefaultCoverageFileProvider" RUNNER="coverage.py" COVERAGE_BY_TEST_ENABLED="false" COVERAGE_TRACING_ENABLED="false" WORKING_DIRECTORY="$PROJECT_DIR$/MNIST_Classification" />
+  </component>
+</project>

SPC-UQ/Cubic_Regression/ConformalRegression.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+
+
+class ConformalRegressionNet(nn.Module):
+    """
+    Simple feedforward regression model with dropout.
+    Output: point prediction only (no uncertainty head).
+    """
+    def __init__(self, input_dim=1, hidden_dim=64, output_dim=1):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.out = nn.Linear(hidden_dim, output_dim)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(p=0.2)
+
+    def forward(self, x):
+        x = self.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.relu(self.fc2(x))
+        x = self.dropout(x)
+        return self.out(x)
+
+
+class ConformalRegressor:
+    """
+    Quantile-based conformal prediction regression model.
+    """
+    def __init__(self, quantile=0.9, learning_rate=5e-3):
+        torch.manual_seed(24)
+        self.quantile = quantile
+        self.model = ConformalRegressionNet()
+        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
+        self.criterion = nn.MSELoss()
+        self.quantile_up = 0.0
+        self.quantile_down = 0.0
+
+    def train(self, x, y, num_epochs=5000):
+        self.model.train()
+        for epoch in range(num_epochs):
+            self.optimizer.zero_grad()
+            pred = self.model(x)
+            loss = self.criterion(pred, y)
+            loss.backward()
+            self.optimizer.step()
+            if (epoch + 1) % 100 == 0:
+                print(f"[Epoch {epoch + 1}/{num_epochs}] Loss: {loss.item():.4f}")
+
+    def calibrate(self, x_calib, y_calib):
+        """
+        Compute empirical quantiles from residuals on calibration set.
+        """
+        self.model.eval()
+        with torch.no_grad():
+            pred = self.model(x_calib).detach().cpu().numpy().squeeze()
+            y_calib_np = y_calib.detach().cpu().numpy().squeeze()
+            residuals = y_calib_np - pred
+
+            res_up = residuals[residuals > 0]
+            res_down = -residuals[residuals <= 0]
+
+            self.quantile_up = np.quantile(res_up, self.quantile) if len(res_up) > 0 else 0.0
+            self.quantile_down = np.quantile(res_down, self.quantile) if len(res_down) > 0 else 0.0
+
+    def predict(self, x):
+        self.model.eval()
+        with torch.no_grad():
+            y_pred = self.model(x).detach().cpu().numpy().squeeze()
+            upper = y_pred + self.quantile_up
+            lower = y_pred - self.quantile_down
+            uncertainty = np.zeros_like(y_pred)  # Conformal prediction doesn't model epistemic
+
+        return y_pred, upper, lower, uncertainty

SPC-UQ/Cubic_Regression/DeepEnsembleRegression.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+
+def nll_loss(mean, log_var, target):
+    """
+    Negative log-likelihood loss for Gaussian output
+    """
+    var = torch.exp(log_var)
+    loss = 0.5 * torch.log(2 * np.pi * var) + 0.5 * ((target - mean) ** 2) / var
+    return loss.mean()
+
+class NLLRegressionNN(nn.Module):
+    """
+    Neural network for regression with Gaussian likelihood output
+    Outputs mean and log-variance
+    """
+    def __init__(self, input_dim=1, hidden_dim=64):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.hidden = nn.Linear(hidden_dim, hidden_dim)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_dim, 2)  # Outputs: mean and log_variance
+
+    def forward(self, x):
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.hidden(x))
+        x = self.fc2(x)
+        mean = x[:, :1]
+        log_var = x[:, 1:]
+        return mean, log_var
+
+class DeepEnsemble:
+    """
+    Deep Ensemble for probabilistic regression with Gaussian likelihood
+    """
+    def __init__(self, num_models=5, learning_rate=5e-3):
+        torch.manual_seed(42)
+        self.models = [NLLRegressionNN() for _ in range(num_models)]
+        self.optimizers = [optim.Adam(model.parameters(), lr=learning_rate) for model in self.models]
+
+    def train(self, data, target, num_epochs=5000):
+        """
+        Train all models independently on the same data
+        """
+        for idx, (model, optimizer) in enumerate(zip(self.models, self.optimizers), start=1):
+            torch.manual_seed(idx + 42)  # Different seed for each model
+            for epoch in range(num_epochs):
+                model.train()
+                optimizer.zero_grad()
+                mean, log_var = model(data)
+                loss = nll_loss(mean, log_var, target)
+                loss.backward()
+                optimizer.step()
+
+                if (epoch + 1) % 500 == 0:
+                    print(f"Model {idx}, Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
+
+    def predict(self, data):
+        """
+        Return ensemble mean, uncertainty, and prediction interval
+        """
+        means = []
+        variances = []
+
+        with torch.no_grad():
+            for model in self.models:
+                model.eval()
+                mean, log_var = model(data)
+                means.append(mean.numpy())
+                variances.append(torch.exp(log_var).numpy())
+
+        means = np.array(means)  # (num_models, batch, 1)
+        variances = np.array(variances)
+
+        # Mean and total predictive variance (mean of model variances)
+        mean_ensemble = np.mean(means, axis=0)
+        var_ensemble = np.mean(variances, axis=0)
+
+        # Epistemic uncertainty: variance across model means
+        epistemic_uncertainty = np.var(means, axis=0)
+
+        std_ensemble = np.sqrt(var_ensemble)
+        y_low = mean_ensemble - 2 * std_ensemble
+        y_high = mean_ensemble + 2 * std_ensemble
+
+        return (
+            mean_ensemble.squeeze(),
+            y_high.squeeze(),
+            y_low.squeeze(),
+            epistemic_uncertainty.squeeze()
+        )

SPC-UQ/Cubic_Regression/EDLQuantileRegression.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import numpy as np
+import matplotlib.pyplot as plt
+import math
+
+
+class DenseNormalGamma(nn.Module):
+    def __init__(self, in_features, out_units):
+        super().__init__()
+        self.out_units = int(out_units)
+        self.linear = nn.Linear(in_features, 4 * self.out_units)
+
+    def evidence(self, x):
+        return F.softplus(x)
+
+    def forward(self, x):
+        output = self.linear(x)
+        mu, log_v, log_alpha, log_beta = torch.chunk(output, chunks=4, dim=-1)
+        v = self.evidence(log_v)
+        alpha = self.evidence(log_alpha) + 1.0
+        beta = self.evidence(log_beta)
+        return torch.cat([mu, v, alpha, beta], dim=-1)
+
+    def extra_repr(self):
+        return f"out_units={self.out_units}"
+
+
+class EDLQRNet(nn.Module):
+    def __init__(self, input_dim=1, num_quantiles=3, hidden_dim=64, num_layers=2, activation=nn.ReLU()):
+        super().__init__()
+        layers = []
+        in_features = input_dim
+        for _ in range(num_layers):
+            layers.append(nn.Linear(in_features, hidden_dim))
+            layers.append(activation)
+            in_features = hidden_dim
+        layers.append(DenseNormalGamma(in_features, num_quantiles))
+        self.network = nn.Sequential(*layers)
+
+    def forward(self, x):
+        output = self.network(x)
+        mu, v, alpha, beta = torch.chunk(output, 4, dim=-1)
+        return mu, v, alpha, beta
+
+
+def nig_nll(y, mu, v, alpha, beta, wi_mean, quantile, reduce=True):
+    tau2 = 2.0 / (quantile * (1.0 - quantile))
+    two_b_lambda = 4.0 * beta * (1.0 + tau2 * wi_mean * v)
+
+    nll = 0.5 * torch.log(math.pi / v) \
+        - alpha * torch.log(two_b_lambda) \
+        + (alpha + 0.5) * torch.log(v * (y - mu) ** 2 + two_b_lambda) \
+        + torch.lgamma(alpha) \
+        - torch.lgamma(alpha + 0.5)
+
+    return torch.mean(nll) if reduce else nll
+
+
+def kl_nig(mu1, v1, a1, b1, mu2, v2, a2, b2):
+    kl = 0.5 * (a1 - 1.0) / b1 * (v2 * (mu2 - mu1) ** 2) \
+        + 0.5 * (v2 / v1) \
+        - 0.5 * torch.log(torch.abs(v2) / torch.abs(v1)) \
+        - 0.5 \
+        + a2 * torch.log(b1 / b2) \
+        - (torch.lgamma(a1) - torch.lgamma(a2)) \
+        + (a1 - a2) * torch.digamma(a1) \
+        - (b1 - b2) * a1 / b1
+    return kl
+
+
+def tilted_loss(q, e):
+    return torch.maximum(q * e, (q - 1.0) * e)
+
+
+def nig_regularization(y, mu, v, alpha, beta, wi_mean, quantile, lambda_reg=0.01, reduce=True, use_kl=False):
+    theta = (1.0 - 2.0 * quantile) / (quantile * (1.0 - quantile))
+    error = tilted_loss(quantile, y - mu)
+
+    if use_kl:
+        kl_val = kl_nig(mu, v, alpha, beta, mu, lambda_reg, 1.0 + lambda_reg, beta)
+        reg = error * kl_val
+    else:
+        evidential_term = 2.0 * v + alpha + 1.0 / beta
+        reg = error * evidential_term
+
+    return torch.mean(reg) if reduce else reg
+
+
+def quantile_evidential_loss(y_true, mu, v, alpha, beta, quantile, coeff=1.0, reduce=True):
+    theta = (1.0 - 2.0 * quantile) / (quantile * (1.0 - quantile))
+    wi_mean = beta / (alpha - 1.0)
+    mu_adj = mu + theta * wi_mean
+
+    loss_nll = nig_nll(y_true, mu_adj, v, alpha, beta, wi_mean, quantile, reduce)
+    loss_reg = nig_regularization(y_true, mu, v, alpha, beta, wi_mean, quantile, reduce)
+    return loss_nll + coeff * loss_reg
+
+
+class EDLQuantileRegressor:
+    def __init__(self, tau_low=0.05, tau_high=0.95, learning_rate=5e-4):
+        torch.manual_seed(42)
+        self.model = EDLQRNet()
+        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
+        self.quantiles = [tau_low, 0.5, tau_high]
+        self.coeff = 0.05
+
+    def loss_function(self, y, mu, v, alpha, beta):
+        total_loss = 0.0
+        for i, q in enumerate(self.quantiles):
+            total_loss += quantile_evidential_loss(
+                y, mu[:, i].unsqueeze(1), v[:, i].unsqueeze(1),
+                alpha[:, i].unsqueeze(1), beta[:, i].unsqueeze(1),
+                q, coeff=self.coeff
+            )
+        return total_loss
+
+    def train(self, x, y, num_epochs=5000):
+        self.model.train()
+        for epoch in range(num_epochs):
+            self.optimizer.zero_grad()
+            mu, v, alpha, beta = self.model(x)
+            loss = self.loss_function(y, mu, v, alpha, beta)
+            loss.backward()
+            self.optimizer.step()
+
+            if (epoch + 1) % 100 == 0:
+                print(f"Epoch [{epoch + 1}/{num_epochs}] Loss: {loss.item():.4f}")
+
+    def predict(self, x):
+        self.model.eval()
+        with torch.no_grad():
+            mu, v, alpha, beta = self.model(x)
+            mu_low, mu_mid, mu_high = torch.unbind(mu, dim=1)
+            v_low, v_mid, v_high = torch.unbind(v, dim=1)
+            alpha_low, alpha_mid, alpha_high = torch.unbind(alpha, dim=1)
+            beta_low, beta_mid, beta_high = torch.unbind(beta, dim=1)
+
+            aleatoric = beta_mid / (alpha_mid - 1.0)
+            epistemic_mid = beta_mid / (v_mid * (alpha_mid - 1.0))
+            epistemic_low = beta_low / (v_low * (alpha_low - 1.0))
+            epistemic_high = beta_high / (v_high * (alpha_high - 1.0))
+            uncertainty = epistemic_mid
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(x, epistemic_mid, label='Mid', color='orange')
+        plt.plot(x, epistemic_low, label='Low', color='red')
+        plt.plot(x, epistemic_high, label='High', color='green')
+        plt.legend()
+        plt.title('Epistemic Uncertainty')
+        plt.show()
+
+        return mu_mid.numpy().squeeze(), mu_high.numpy().squeeze(), mu_low.numpy().squeeze(), uncertainty.numpy().squeeze()

SPC-UQ/Cubic_Regression/EDLRegression.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from scipy.stats import t as student_t
+import numpy as np
+import math
+
+
+class EDLRegressionNet(nn.Module):
+    def __init__(self, input_dim=1, hidden_dim=64, output_dim=1):
+        super().__init__()
+        self.hidden1 = nn.Linear(input_dim, hidden_dim)
+        self.hidden2 = nn.Linear(hidden_dim, hidden_dim)
+        self.output_mu = nn.Linear(hidden_dim, output_dim)
+        self.output_logv = nn.Linear(hidden_dim, output_dim)
+        self.output_alpha = nn.Linear(hidden_dim, output_dim)
+        self.output_beta = nn.Linear(hidden_dim, output_dim)
+        self.activation = nn.ReLU()
+
+    def forward(self, x):
+        x = self.activation(self.hidden1(x))
+        x = self.activation(self.hidden2(x))
+        mu = self.output_mu(x)
+        v = F.softplus(self.output_logv(x))
+        alpha = F.softplus(self.output_alpha(x)) + 1.0 + 1e-6
+        beta = F.softplus(self.output_beta(x))
+        return mu, v, alpha, beta
+
+
+def nig_nll(y, mu, v, alpha, beta, reduce=True):
+    two_b_lambda = 2 * beta * (1 + v)
+
+    nll = 0.5 * torch.log(torch.tensor(np.pi) / v) \
+        - alpha * torch.log(two_b_lambda) \
+        + (alpha + 0.5) * torch.log(v * (y - mu) ** 2 + two_b_lambda) \
+        + torch.lgamma(alpha) \
+        - torch.lgamma(alpha + 0.5)
+
+    return torch.mean(nll) if reduce else nll
+
+
+def kl_nig(mu1, v1, a1, b1, mu2, v2, a2, b2):
+    eps = 1e-6
+    v1 = torch.clamp(v1, min=eps)
+    v2 = torch.clamp(v2, min=eps)
+    b1 = torch.clamp(b1, min=eps)
+    b2 = torch.clamp(b2, min=eps)
+
+    term1 = 0.5 * (a1 - 1) / b1 * (v2 * (mu2 - mu1) ** 2)
+    term2 = 0.5 * v2 / v1
+    term3 = -0.5 * torch.log(v2 / v1)
+    term4 = -0.5
+    term5 = a2 * torch.log(b1 / b2)
+    term6 = -(torch.lgamma(a1) - torch.lgamma(a2))
+    term7 = (a1 - a2) * torch.digamma(a1)
+    term8 = -(b1 - b2) * a1 / b1
+
+    kl = term1 + term2 + term3 + term4 + term5 + term6 + term7 + term8
+    return kl
+
+
+def nig_regularization(y, mu, v, alpha, beta, omega=0.01, reduce=True, use_kl=False):
+    error = torch.abs(y - mu)
+
+    if use_kl:
+        kl = kl_nig(mu, v, alpha, beta, mu, omega, 1 + omega, beta)
+        reg = error * kl
+    else:
+        evidential = 2 * v + alpha
+        reg = error * evidential
+
+    return reg.mean() if reduce else reg
+
+
+def edl_loss(y, mu, v, alpha, beta, lam=0.0, reduce=True, return_components=False):
+    nll = nig_nll(y, mu, v, alpha, beta, reduce=reduce)
+    reg = nig_regularization(y, mu, v, alpha, beta, reduce=reduce)
+    loss = nll  # optionally: loss = nll + lam * reg
+
+    return (loss, (nll, reg)) if return_components else loss
+
+
+def predictive_interval(mu, v, alpha, beta, confidence=0.95):
+    mu = torch.as_tensor(mu)
+    v = torch.as_tensor(v)
+    alpha = torch.as_tensor(alpha)
+    beta = torch.as_tensor(beta)
+
+    dof = 2.0 * alpha
+    scale = torch.sqrt((1.0 + v) * beta / (alpha * v))
+    lower_q = (1.0 - confidence) / 2.0
+    upper_q = 1.0 - lower_q
+
+    t_l = student_t.ppf(lower_q, df=dof.cpu().numpy())
+    t_u = student_t.ppf(upper_q, df=dof.cpu().numpy())
+
+    lower = mu + torch.from_numpy(t_l).to(mu.device) * scale
+    upper = mu + torch.from_numpy(t_u).to(mu.device) * scale
+
+    return lower, upper
+
+
+class EDLRegressor:
+    def __init__(self, learning_rate=5e-4):
+        torch.manual_seed(33)
+        self.model = EDLRegressionNet()
+        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
+        self.criterion = edl_loss
+        self.lambda_ = 0.01
+
+    def train(self, x, y, num_epochs=5000):
+        torch.manual_seed(33)
+        self.model.train()
+        for epoch in range(num_epochs):
+            self.optimizer.zero_grad()
+            mu, v, alpha, beta = self.model(x)
+            loss = self.criterion(y, mu, v, alpha, beta, lam=self.lambda_)
+            loss.backward()
+            self.optimizer.step()
+            if (epoch + 1) % 100 == 0:
+                print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
+
+    def predict(self, x):
+        self.model.eval()
+        with torch.no_grad():
+            mu, v, alpha, beta = self.model(x)
+
+            # Uncertainty decomposition
+            aleatoric = torch.sqrt(beta / (alpha - 1.0 + 1e-6))
+            epistemic = torch.sqrt(beta / (v * (alpha - 1.0 + 1e-6)))
+
+            mu_np = mu.detach().cpu().numpy()
+            aleatoric_np = aleatoric.detach().cpu().numpy()
+            epistemic_np = epistemic.detach().cpu().numpy()
+
+            lower, upper = predictive_interval(mu, v, alpha, beta, confidence=0.95)
+            lower_np = lower.detach().cpu().numpy()
+            upper_np = upper.detach().cpu().numpy()
+
+        return mu_np.squeeze(), upper_np.squeeze(), lower_np.squeeze(), epistemic_np.squeeze()

SPC-UQ/Cubic_Regression/QROC.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+from torch.utils.data import DataLoader, TensorDataset
+
+def pinball_loss(q_pred, target, tau):
+    """
+    Compute the quantile (pinball) loss for a given quantile level tau
+    """
+    error = target - q_pred
+    return torch.mean(torch.max(tau * error, (tau - 1) * error))
+
+def multi_quantile_loss(q_low, q_mid, q_high, target, tau_low=0.05, tau_high=0.95):
+    """
+    Compute total loss across low, median, and high quantiles
+    """
+    loss_l = pinball_loss(q_low, target, tau_low)
+    loss_m = pinball_loss(q_mid, target, 0.5)
+    loss_h = pinball_loss(q_high, target, tau_high)
+    return loss_l + loss_m + loss_h
+
+
+class QuantileRegressionNN(nn.Module):
+    """
+    Fully-connected quantile regression network with three outputs:
+    lower, median, upper quantiles.
+    """
+    def __init__(self, input_dim=1, hidden_dim=64):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.out = nn.Linear(hidden_dim, 3)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.fc2(x))
+        out = self.out(x)
+        return out[:, :1], out[:, 1:2], out[:, 2:]
+
+    def extract_features(self, x):
+        """
+        Extract penultimate-layer features (for certificate head)
+        """
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.fc2(x))
+        return x
+
+
+def build_certificate_head(features, out_dim=20, epochs=500):
+    """
+    Train an orthogonal linear projection (certificate head) on extracted features
+    """
+    projection = nn.Linear(features.size(1), out_dim)
+    loader = DataLoader(TensorDataset(features), shuffle=True, batch_size=128)
+    optimizer = optim.Adam(projection.parameters())
+
+    for _ in range(epochs):
+        for (feature_batch,) in loader:
+            optimizer.zero_grad()
+            output = projection(feature_batch)
+            cert_loss = output.pow(2).mean()
+
+            identity = torch.eye(out_dim, device=projection.weight.device)
+            ortho_penalty = (projection.weight @ projection.weight.T - identity).pow(2).mean()
+
+            (cert_loss + ortho_penalty).backward()
+            optimizer.step()
+
+    return projection
+
+
+class QROC:
+    """
+    Single-model Quantile Regression with Orthogonal Certificate head (QROC)
+    Provides aleatoric and epistemic uncertainty estimation
+    """
+    def __init__(self, learning_rate=5e-3, tau_low=0.05, tau_high=0.95):
+        torch.manual_seed(42)
+        self.model = QuantileRegressionNN()
+        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
+        self.tau_low = tau_low
+        self.tau_high = tau_high
+        self.certificate_head = None
+
+    def train(self, x, y, epochs=3000):
+        """
+        Train the quantile regression network and then fit certificate head on extracted features
+        """
+        for epoch in range(epochs):
+            self.model.train()
+            self.optimizer.zero_grad()
+            q_low, q_mid, q_high = self.model(x)
+            loss = multi_quantile_loss(q_low, q_mid, q_high, y, self.tau_low, self.tau_high)
+            loss.backward()
+            self.optimizer.step()
+
+            if (epoch + 1) % 500 == 0:
+                print(f"[{epoch+1}/{epochs}] Loss: {loss.item():.4f}")
+
+        # Train certificate head using the final feature representation
+        self.model.eval()
+        with torch.no_grad():
+            features = self.model.extract_features(x).detach()
+        self.certificate_head = build_certificate_head(features)
+
+    def predict(self, x):
+        """
+        Predict quantiles and return aleatoric and epistemic uncertainties
+        """
+        self.model.eval()
+        with torch.no_grad():
+            q_low, q_mid, q_high = self.model(x)
+
+            # Epistemic: projection energy via orthogonal certificate head
+            features = self.model.extract_features(x)
+            epistemic = self.certificate_head(features).pow(2).mean(dim=1).cpu().numpy()
+
+        return (
+            q_mid.squeeze().numpy(),
+            q_high.squeeze().numpy(),
+            q_low.squeeze().numpy(),
+            epistemic
+        )

SPC-UQ/Cubic_Regression/SPCRegression.py

@@ -0,0 +1,173 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+def pinball_loss(q_pred, target, tau):
+    """Standard quantile regression loss."""
+    errors = target - q_pred
+    loss = torch.max(tau * errors, (tau - 1) * errors)
+    return torch.mean(loss)
+
+def cali_loss(y_pred, y_true, q, scale=True):
+    """
+    Calibration loss for quantile regression.
+    Penalizes over- or under-coverage relative to quantile level q.
+    """
+    diff = y_true - y_pred
+    under_mask = (y_true <= y_pred)
+    over_mask = ~under_mask
+
+    coverage = torch.mean(under_mask.float())
+
+    if coverage < q:
+        loss = torch.mean(diff[over_mask])
+    else:
+        loss = torch.mean(-diff[under_mask])
+
+    if scale:
+        loss *= torch.abs(q - coverage)
+    return loss
+
+class SPCRegressionNet(nn.Module):
+    """
+    Neural network that predicts:
+    - point estimate (v)
+    - MAR (mean absolute residual)
+    - MAR up/down (for epistemic decomposition)
+    - QR up/down (for aleatoric decomposition)
+    """
+    def __init__(self, input_dim=1, hidden_dim=64):
+        super(SPCRegressionNet, self).__init__()
+        self.hidden = nn.Linear(input_dim, hidden_dim)
+        self.hidden2 = nn.Linear(hidden_dim, hidden_dim)
+        self.hidden3 = nn.Linear(hidden_dim, hidden_dim)
+
+        self.relu = nn.ReLU()
+        # self.dropout = nn.Dropout(p=0.2)
+        self.output_v = nn.Linear(hidden_dim, 1)
+        self.output_uq = nn.Linear(hidden_dim, 5)
+
+
+    def forward(self, x):
+        x = self.hidden(x)
+        x = self.relu(x)
+        x = self.hidden2(x)
+        x = self.relu(x)
+        v = self.output_v(x)
+        x = self.hidden3(x)
+        x = self.relu(x)
+        output = self.output_uq(x)
+        mar, mar_up, mar_down, q_up, q_down = torch.chunk(output, 5, dim=-1)
+
+        q_up = F.softplus(q_up)
+        q_down = F.softplus(q_down)
+        return v, mar, mar_up, mar_down, q_up, q_down
+
+class SPCregression:
+    """
+    Trainer and predictor for SPC UQ model.
+    Supports joint or stagewise training strategies.
+    """
+    def __init__(self, learning_rate=5e-3):
+        torch.manual_seed(42)
+        self.model = SPCRegressionNet()
+        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
+        self.optimizer1 = optim.Adam(list(self.model.hidden.parameters()) + list(self.model.hidden2.parameters()) + list(self.model.output_v.parameters()),
+                                     lr=learning_rate,weight_decay=1e-4)
+        self.optimizer2 = optim.Adam(list(self.model.hidden3.parameters()) + list(self.model.output_uq.parameters()),
+                                     lr=learning_rate,weight_decay=1e-4)
+        self.criterion = nn.MSELoss()
+        self.criterion2=nn.L1Loss()
+
+    def mar_loss(self, y, predictions, mar, mar_up, mar_down, q_up, q_down):
+        """Computes loss for MAR and QR heads."""
+        residual = abs(y - predictions)
+        diff = (y - predictions.detach())
+        loss_mar = self.criterion(mar, residual)
+
+        mask_up = (diff > 0)
+        mask_down = (diff < 0)
+        loss_mar_up = self.criterion(mar_up[mask_up], (y[mask_up] - predictions[mask_up]))
+        loss_mar_down = self.criterion(mar_down[mask_down], (predictions[mask_down] - y[mask_down]))
+
+        # loss_q_up = pinball_loss(q_up[mask_up], (y[mask_up] - predictions[mask_up]), 0.95)
+        # loss_q_down = pinball_loss(q_down[mask_down], (predictions[mask_down] - y[mask_down]), 0.95)
+        loss_cali_up = cali_loss(q_up[mask_up], (y[mask_up] - predictions[mask_up]), 0.95)
+        loss_cali_down = cali_loss(q_down[mask_down], (predictions[mask_down] - y[mask_down]), 0.95)
+
+
+        loss = loss_mar + loss_mar_up + loss_mar_down +(loss_cali_up + loss_cali_down)
+               # + 0.2 * (loss_q_up + loss_q_down) \
+               # + 0.8 * (loss_cali_up + loss_cali_down) \
+        return loss
+
+    def train(self, data, target, num_epochs=5000, strategy='stagewise'):
+        """
+        Train the model using either:
+        - 'joint': full loss on all components
+        - 'stagewise': first fit task head, then UQ heads
+        """
+        torch.manual_seed(42)
+        self.model.train()
+
+        if strategy== 'joint':
+            for epoch in range(num_epochs):
+                self.optimizer.zero_grad()
+                predictions, mar, mar_up, mar_down, q_up, q_down = self.model(data)
+                loss = self.criterion(predictions, target)+self.mar_loss(target, predictions, mar, mar_up, mar_down, q_up, q_down)
+                loss.backward()
+                self.optimizer.step()
+                if (epoch + 1) % 100 == 0:
+                    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
+
+        if strategy== 'stagewise':
+            for epoch in range(num_epochs):
+                self.optimizer1.zero_grad()
+                predictions, mar, mar_up, mar_down, q_up, q_down = self.model(data)
+                loss = self.criterion(predictions, target)
+                loss.backward()
+                self.optimizer1.step()
+                if (epoch + 1) % 100 == 0:
+                    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
+
+            for epoch in range(num_epochs):
+                self.optimizer2.zero_grad()
+                predictions, mar, mar_up, mar_down, q_up, q_down = self.model(data)
+                loss = self.mar_loss(target, predictions, mar, mar_up, mar_down, q_up, q_down)
+                loss.backward()
+                self.optimizer2.step()
+                if (epoch + 1) % 100 == 0:
+                    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")
+
+    def predict(self, data, calibration=False):
+        """Run prediction and return interval bounds and uncertainty estimate."""
+        self.model.eval()
+        with torch.no_grad():
+            predictions, mar, mar_up, mar_down, q_up, q_down = self.model(data)
+
+            v = predictions.numpy()
+            mar = mar.detach().numpy()
+            mar_up = mar_up.detach().numpy()
+            mar_down = mar_down.detach().numpy()
+            q_up = q_up.detach().numpy()
+            q_down = q_down.detach().numpy()
+
+            if calibration:
+                # Calibration adjustment based on Self-consistency
+                d_up = (mar * mar_down) / ((2 * mar_down - mar) * mar_up)
+                d_down = (mar * mar_up) / ((2 * mar_up - mar) * mar_down)
+                d_up = np.clip(d_up, 1, None)
+                d_down = np.clip(d_down, 1, None)
+                q_up *= d_up
+                q_down *= d_down
+
+            high_bound = v + q_up
+            low_bound = v - q_down
+
+            # Self-consistency Verification
+            uncertainty = (abs(2 * mar_up * mar_down - mar * (mar_up + mar_down)))
+
+
+        return v.squeeze(), high_bound.squeeze(), low_bound.squeeze(), uncertainty.squeeze()

SPC-UQ/Cubic_Regression/run_cubic_tests.py

@@ -0,0 +1,335 @@
+import torch
+import numpy as np
+import argparse
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+from SPCRegression import SPCregression
+from DeepEnsembleRegression import DeepEnsemble
+from ConformalRegression import ConformalRegressor
+from EDLRegression import EDLRegressor
+from EDLQuantileRegression import EDLQuantileRegressor
+from QROC import QROC
+from scipy.stats import binned_statistic
+
+
+def ece_pi(y_true, pred_lower, pred_upper, num_bins=10):
+    """Compute Expected Calibration Error (ECE) based on prediction interval width bins."""
+    N = y_true.shape[0]
+    in_interval = ((y_true >= pred_lower) & (y_true <= pred_upper)).astype(float)
+    widths = pred_upper - pred_lower
+    min_w, max_w = np.min(widths), np.max(widths)
+
+    if min_w == max_w:
+        return np.abs(in_interval.mean() - 1.0)
+
+    bin_edges = np.linspace(min_w, max_w, num_bins + 1)
+    ece = 0.0
+
+    for i in range(num_bins):
+        bin_mask = (widths >= bin_edges[i]) & (widths < bin_edges[i + 1])
+        count_in_bin = np.sum(bin_mask)
+        if count_in_bin == 0:
+            continue
+        avg_in_interval = in_interval[bin_mask].mean()
+        nominal_coverage = 0.95
+        calib_error = np.abs(avg_in_interval - nominal_coverage)
+        weight = count_in_bin / N
+        ece += weight * calib_error
+    return ece
+
+
+def binning(pred_lower, pred_upper, num_bins=10):
+    """Group indices into bins based on interval width."""
+    widths = pred_upper - pred_lower
+    min_w, max_w = np.min(widths), np.max(widths)
+    bin_edges = np.linspace(min_w, max_w, num_bins + 1)
+
+    bins = []
+    for i in range(num_bins):
+        if i == num_bins - 1:
+            bin_mask = (widths >= bin_edges[i]) & (widths <= bin_edges[i + 1])
+        else:
+            bin_mask = (widths >= bin_edges[i]) & (widths < bin_edges[i + 1])
+        bin_indices = np.where(bin_mask)[0]
+        bins.append(bin_indices)
+    return bins, bin_edges
+
+
+def generate_multimodal_data(n_samples=1000):
+    """Generate mixture-distribution noise samples."""
+    x = np.random.randn(n_samples)
+    mask = np.random.choice([0, 1, 2], p=[0.4, 0.3, 0.3], size=n_samples)
+    y = np.where(mask == 0, x + np.random.randn(n_samples) * 1,
+                 np.where(mask == 1, x + 40 + np.random.randn(n_samples) * 1,
+                          x - 10 + np.random.randn(n_samples) * 1))
+    return y
+
+
+def generate_train_data(n_samples=20, noise='log'):
+    """Generate synthetic training data with nonlinear relationship and optional noise."""
+    np.random.seed(57)
+    x = np.linspace(-4, 4, n_samples)
+
+    if noise == 'log':
+        noise = np.random.lognormal(mean=1.5, sigma=1, size=n_samples)
+    elif noise == 'tri':
+        noise = generate_multimodal_data(n_samples)
+    elif noise == 'norm':
+        noise = np.random.normal(0, 8, size=n_samples)
+
+    noise = noise - np.mean(noise)
+    y = x ** 3 + noise
+    x = x.reshape(-1, 1).astype(np.float32)
+    y = y.reshape(-1, 1).astype(np.float32)
+    return x, y
+
+
+def generate_test_data(n_samples=100, noise='log'):
+    """Generate synthetic test data with extended input range and optional noise."""
+    np.random.seed(27)
+    x = np.linspace(-6, 6, n_samples)
+
+    if noise == 'log':
+        noise = np.random.lognormal(mean=1.5, sigma=1, size=n_samples)
+    elif noise == 'tri':
+        noise = generate_multimodal_data(n_samples)
+    elif noise == 'norm':
+        noise = np.random.normal(0, 8, size=n_samples)
+
+    noise = noise - np.mean(noise)
+    y = x ** 3 + noise
+    x = x.reshape(-1, 1).astype(np.float32)
+    y = y.reshape(-1, 1).astype(np.float32)
+    return x, y
+
+
+# ===================== Argument Parser ===================== #
+parser = argparse.ArgumentParser()
+parser.add_argument("--num-epochs", default=5000, type=int)
+parser.add_argument('--data-noise', default='log', choices=['norm', 'tri', 'log'])
+parser.add_argument('--UQ-model', default='SPCregression', choices=['SPCregression', 'DeepEnsemble', 'EDLRegressor', 'EDLQuantileRegressor', 'QROC', 'ConformalRegressor'], help='Select UQ model to test.')
+args = parser.parse_args()
+
+# Generate training, calibration, and testing datasets
+x_train, y_train = generate_train_data(n_samples=2000, noise=args.data_noise)
+x_calib, y_calib = generate_train_data(n_samples=500, noise=args.data_noise)
+x_test, y_test = generate_test_data(n_samples=1000, noise=args.data_noise)
+
+# Convert data to torch tensors
+x_train_tensor = torch.from_numpy(x_train)
+y_train_tensor = torch.from_numpy(y_train)
+x_calib_tensor = torch.from_numpy(x_calib)
+y_calib_tensor = torch.from_numpy(y_calib)
+x_test_tensor = torch.from_numpy(x_test)
+y_test_tensor = torch.from_numpy(y_test)
+
+# Training parameters
+num_epochs = 5000
+num_models = 5
+lr = 0.001
+
+# Select UQ model to test
+UQ = args.UQ_model
+
+if UQ == 'SPCregression':
+    model = SPCregression(learning_rate=lr)
+elif UQ == 'DeepEnsemble':
+    model = DeepEnsemble(num_models=num_models, learning_rate=lr)
+elif UQ == 'EDLRegressor':
+    model = EDLRegressor(learning_rate=lr)
+elif UQ == 'EDLQuantileRegressor':
+    model = EDLQuantileRegressor(tau_low=0.025, tau_high=0.975, learning_rate=lr)
+elif UQ == 'QROC':
+    model = QROC(tau_low=0.025, tau_high=0.975, learning_rate=lr)
+elif UQ == 'ConformalRegressor':
+    model = ConformalRegressor(0.95, learning_rate=lr)
+
+# Train model
+model.train(x_train_tensor, y_train_tensor, num_epochs)
+
+# Calibrate if applicable
+if UQ == 'ConformalRegressor':
+    model.calibrate(x_calib_tensor, y_calib_tensor)
+
+# Predict on train and test sets
+mean, upper_bound, lower_bound, uncertainty = model.predict(x_train_tensor)
+y_train_np = y_train.flatten()
+
+# Evaluate train metrics
+mpi_width = np.mean(upper_bound - lower_bound)
+picp = np.mean(((y_train_np >= lower_bound) & (y_train_np <= upper_bound)).astype(float))
+rmse = np.sqrt(np.mean((y_train_np - mean) ** 2))
+print(f"Train Mean Prediction Interval Width (MPIW): {mpi_width:.4f}")
+print(f"Train Prediction Interval Coverage Probability (PICP): {picp:.4f}")
+print(f"Train Root Mean Squared Error (RMSE): {rmse:.4f}")
+
+# Predict on test set
+mean, upper_bound, lower_bound, uncertainty = model.predict(x_test_tensor)
+x_test_np = x_test.flatten()
+y_test_np = y_test.flatten()
+
+# In-distribution test subset for interval evaluation
+y_low_id = lower_bound[170:830]
+y_high_id = upper_bound[170:830]
+y_mean_id = mean[170:830]
+x_test_id = x_test_np[170:830]
+y_test_id = y_test_np[170:830]
+
+mpi_width_id = np.mean(y_high_id - y_low_id)
+picp_id = np.mean(((y_test_id >= y_low_id) & (y_test_id <= y_high_id)).astype(float))
+# PICP+
+picp_plus = np.sum(((y_test_id >= y_mean_id) & (y_test_id <= y_high_id)).astype(float))/np.sum((y_test_id >= y_mean_id).astype(float))
+# PICP-
+picp_minus = np.sum(((y_test_id >= y_low_id) & (y_test_id <= y_mean_id)).astype(float))/np.sum((y_test_id <= y_mean_id).astype(float))
+
+rmse_id = np.sqrt(np.mean((y_test_id - y_mean_id) ** 2))
+
+print(f"ID Mean Prediction Interval Width (MPIW): {mpi_width_id:.4f}")
+print(f"ID Prediction Interval Coverage Probability (PICP): {picp_id:.4f}")
+print(f"ID Prediction Interval Coverage Probability (PICP+): {picp_plus:.4f}")
+print(f"ID Prediction Interval Coverage Probability (PICP-): {picp_minus:.4f}")
+print(f"ID Root Mean Squared Error (RMSE): {rmse_id:.4f}")
+
+# Compute ECE
+ece = ece_pi(y_test_id, y_low_id, y_high_id, num_bins=10)
+print(f"Expected Calibration Error (ECE): {ece:.4f}")
+
+threshold = uncertainty[170:830].mean()
+print('threshold:', threshold)
+cer_count = 0
+unc_count = 0
+cer_diff = []
+unc_diff = []
+
+for i in range(len(uncertainty)):
+    unc = uncertainty[i]
+    diff_sq = (y_test_np[i] - mean[i])**2
+    if unc < threshold:
+        cer_count += 1
+        cer_diff.append(diff_sq)
+    else:
+        unc_count += 1
+        unc_diff.append(diff_sq)
+
+rmse_certain = np.sqrt(np.mean(cer_diff)) if len(cer_diff)>0 else 0
+rmse_uncertain = np.sqrt(np.mean(unc_diff)) if len(unc_diff)>0 else 0
+rmse_all = np.sqrt(np.mean(cer_diff + unc_diff))
+
+print('Certain sample:', cer_count,
+      'RMSE_certain:', round(rmse_certain,4),
+      'Uncertain sample:', unc_count,
+      'RMSE_uncertain:', round(rmse_uncertain,4),
+      'RMSE_all:', round(rmse_all,4))
+
+print(round(rmse_id,4),
+      round(picp_id,4),
+      round(mpi_width_id,4),
+      round(rmse_certain,4),
+      round(rmse_uncertain,4),
+      round(rmse_all,4),
+      cer_count, unc_count)
+
+
+def plot_binned_intervals(x_test, y_test, mean, lower_bound, upper_bound, bins, num_bins=1):
+    """Plot prediction intervals with color-coded bins based on interval width."""
+    x = x_test.squeeze()
+    y = y_test.squeeze()
+
+    def quantile_stat(q):
+        def func(y_in_bin):
+            return np.percentile(y_in_bin, q)
+        return func
+
+    bin_means, bin_edges, _ = binned_statistic(x, y, statistic='mean', bins=num_bins)
+    q5, _, _ = binned_statistic(x, y, statistic=quantile_stat(5), bins=bin_edges)
+    q95, _, _ = binned_statistic(x, y, statistic=quantile_stat(95), bins=bin_edges)
+
+    gt = x ** 3
+    y_up = (y - gt)[y > gt]
+    y_down = (gt - y)[y < gt]
+    gt_up = np.quantile(y_up, 0.95)
+    gt_down = np.quantile(y_down, 0.95)
+
+    plt.figure(figsize=(7, 5))
+    cmap = cm.get_cmap('viridis', len(bins))
+    interval_width = 0.1
+    for i, bin_indices in enumerate(bins):
+        color = cmap(i)
+        for j, idx in enumerate(bin_indices):
+            x_val = x_test[idx]
+            plt.fill_between(
+                [x_val - interval_width / 2, x_val + interval_width / 2],
+                [lower_bound[idx], lower_bound[idx]],
+                [upper_bound[idx], upper_bound[idx]],
+                color=color,
+                alpha=0.2,
+                label=f'Bin {i + 1}' if j == 0 else None
+            )
+
+    plt.plot(x, gt, color='blue', linestyle='--', label='Mean E[y|x]', linewidth=2)
+    plt.plot(x, gt - gt_down, color='darkorange', linestyle='--', label='Lower bound GT', linewidth=2)
+    plt.plot(x, gt + gt_up, color='purple', linestyle='--', label='Upper bound GT', linewidth=2)
+    plt.scatter(x_test, y_test, color='green', s=5, label='Test Data')
+    plt.plot(x_test, mean, color='red', label='Point Prediction')
+    plt.title("Equal-width Binned PIs")
+    plt.legend()
+    plt.ylim(-100, 100)
+    plt.tight_layout()
+    plt.show()
+
+    # Second plot with shaded split intervals
+    plt.figure(figsize=(7, 4))
+    plt.plot(x, gt, color='blue', linestyle='--', label='Mean E[y|x]', linewidth=2)
+    plt.plot(x, gt - gt_down, color='darkorange', linestyle='--', label='Lower bound GT', linewidth=2)
+    plt.plot(x, gt + gt_up, color='purple', linestyle='--', label='Upper bound GT', linewidth=2)
+    plt.scatter(x_test, y_test, color='green', s=5, label='Test Data')
+    plt.plot(x_test, mean, color='red', label='Point Prediction')
+    plt.fill_between(x_test.flatten(), lower_bound, mean, color='orange', alpha=0.3, label='Lower Interval')
+    plt.fill_between(x_test.flatten(), mean, upper_bound, color='purple', alpha=0.4, label='Upper Interval')
+    plt.legend(fontsize=8)
+    plt.ylim(-80, 100)
+    plt.title('Split-point Prediction Intervals')
+    plt.tight_layout()
+    plt.show()
+
+def plot_intervals(x_test, y_test, mean, lower_bound, upper_bound, uncertainty):
+    """Plot prediction intervals and epistemic uncertainty with ground truth reference."""
+    x = x_test.squeeze()
+    y = y_test.squeeze()
+
+    gt = x ** 3
+    y_up = (y - gt)[y > gt]
+    y_down = (gt - y)[y < gt]
+    gt_up = np.quantile(y_up, 0.95)
+    gt_down = np.quantile(y_down, 0.95)
+
+    import matplotlib.gridspec as gridspec
+    plt.figure(figsize=(7, 9))
+    gs = gridspec.GridSpec(2, 1, height_ratios=[4, 3])
+
+    ax1 = plt.subplot(gs[0])
+    ax1.axvspan(x.min(), -4, facecolor='lightgray', alpha=0.4)
+    ax1.axvspan(4, x.max(), facecolor='lightgray', alpha=0.4)
+    ax1.plot(x, gt, color='blue', linestyle='--', label='Mean E[y|x]', linewidth=2)
+    ax1.plot(x, gt - gt_down, color='darkorange', linestyle='--', label='Lower bound GT', linewidth=2)
+    ax1.plot(x, gt + gt_up, color='purple', linestyle='--', label='Upper bound GT', linewidth=2)
+    ax1.scatter(x_test, y_test, color='green', s=5, label='Test Data')
+    ax1.plot(x_test, mean, color='red', label='Point Prediction')
+    ax1.fill_between(x_test.flatten(), lower_bound, mean, color='orange', alpha=0.3, label='Lower Interval')
+    ax1.fill_between(x_test.flatten(), mean, upper_bound, color='purple', alpha=0.4, label='Upper Interval')
+    ax1.legend(fontsize=9)
+    ax1.set_ylim(-150, 150)
+
+    ax2 = plt.subplot(gs[1])
+    ax2.axvspan(x.min(), -4, facecolor='lightgray', alpha=0.4, label='OOD Region')
+    ax2.axvspan(4, x.max(), facecolor='lightgray', alpha=0.4)
+    ax2.plot(x_test, uncertainty, label='Epistemic uncertainty', color='dodgerblue')
+    ax2.fill_between(x_test.squeeze(), uncertainty.squeeze(), alpha=0.3, color='dodgerblue')
+    ax2.legend()
+    plt.tight_layout()
+    plt.show()
+
+# Call visualizations after metric evaluations
+bins, bin_edges = binning(y_low_id, y_high_id, num_bins=5)
+plot_binned_intervals(x_test_id, y_test_id, y_mean_id, y_low_id, y_high_id, bins)
+plot_intervals(x_test, y_test, mean, lower_bound, upper_bound, uncertainty)

SPC-UQ/Image_Classification/README.md

@@ -0,0 +1,285 @@
+# Deep Deterministic Uncertainty
+
+[![arXiv](https://img.shields.io/badge/stat.ML-arXiv%3A2006.08437-B31B1B.svg)](https://arxiv.org/abs/2102.11582)
+[![Pytorch 1.8.1](https://img.shields.io/badge/pytorch-1.8.1-blue.svg)](https://pytorch.org/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/omegafragger/DDU/blob/main/LICENSE)
+
+This repository contains the code for [*Deterministic Neural Networks with Appropriate Inductive Biases Capture Epistemic and Aleatoric Uncertainty*](https://arxiv.org/abs/2102.11582).
+
+If the code or the paper has been useful in your research, please add a citation to our work:
+
+```
+@article{mukhoti2021deterministic,
+  title={Deterministic Neural Networks with Appropriate Inductive Biases Capture Epistemic and Aleatoric Uncertainty},
+  author={Mukhoti, Jishnu and Kirsch, Andreas and van Amersfoort, Joost and Torr, Philip HS and Gal, Yarin},
+  journal={arXiv preprint arXiv:2102.11582},
+  year={2021}
+}
+```
+
+## Dependencies
+
+The code is based on PyTorch and requires a few further dependencies, listed in [environment.yml](environment.yml). It should work with newer versions as well.
+
+
+## OoD Detection
+
+### Datasets
+
+For OoD detection, you can train on [*CIFAR-10/100*](https://www.cs.toronto.edu/~kriz/cifar.html). You can also train on [*Dirty-MNIST*](https://blackhc.github.io/ddu_dirty_mnist/) by downloading *Ambiguous-MNIST* (```amnist_labels.pt``` and ```amnist_samples.pt```) from [here](https://github.com/BlackHC/ddu_dirty_mnist/releases/tag/data-v0.6.0) and using the following training instructions.
+
+### Training
+
+In order to train a model for the OoD detection task, use the [train.py](train.py) script. Following are the main parameters for training:
+```
+--seed: seed for initialization
+--dataset: dataset used for training (cifar10/cifar100/dirty_mnist)
+--dataset-root: /path/to/amnist_labels.pt and amnist_samples.pt/ (if training on dirty-mnist)
+--model: model to train (wide_resnet/vgg16/resnet18/resnet50/lenet)
+-sn: whether to use spectral normalization (available for wide_resnet, vgg16 and resnets)
+--coeff: Coefficient for spectral normalization
+-mod: whether to use architectural modifications (leaky ReLU + average pooling in skip connections)
+--save-path: path/for/saving/model/
+```
+
+As an example, in order to train a Wide-ResNet-28-10 with spectral normalization and architectural modifications on CIFAR-10, use the following:
+```
+python train.py \
+       --seed 1 \
+       --dataset cifar10 \
+       --model wide_resnet \
+       -sn -mod \
+       --coeff 3.0 
+```
+Similarly, to train a ResNet-18 with spectral normalization on Dirty-MNIST, use:
+```
+python train.py \
+       --seed 1 \
+       --dataset dirty-mnist \
+       --dataset-root /home/user/amnist/ \
+       --model resnet18 \
+       -sn \
+       --coeff 3.0
+```
+
+### Evaluation
+
+To evaluate trained models, use [evaluate.py](evaluate.py). This script can evaluate and aggregate results over multiple experimental runs. For example, if the pretrained models are stored in a directory path ```/home/user/models```, store them using the following directory structure:
+```
+models
+├── Run1
+│   └── wide_resnet_1_350.model
+├── Run2
+│   └── wide_resnet_2_350.model
+├── Run3
+│   └── wide_resnet_3_350.model
+├── Run4
+│   └── wide_resnet_4_350.model
+└── Run5
+    └── wide_resnet_5_350.model
+```
+For an ensemble of models, store the models using the following directory structure:
+```
+model_ensemble
+├── Run1
+│   ├── wide_resnet_1_350.model
+│   ├── wide_resnet_2_350.model
+│   ├── wide_resnet_3_350.model
+│   ├── wide_resnet_4_350.model
+│   └── wide_resnet_5_350.model
+├── Run2
+│   ├── wide_resnet_10_350.model
+│   ├── wide_resnet_6_350.model
+│   ├── wide_resnet_7_350.model
+│   ├── wide_resnet_8_350.model
+│   └── wide_resnet_9_350.model
+├── Run3
+│   ├── wide_resnet_11_350.model
+│   ├── wide_resnet_12_350.model
+│   ├── wide_resnet_13_350.model
+│   ├── wide_resnet_14_350.model
+│   └── wide_resnet_15_350.model
+├── Run4
+│   ├── wide_resnet_16_350.model
+│   ├── wide_resnet_17_350.model
+│   ├── wide_resnet_18_350.model
+│   ├── wide_resnet_19_350.model
+│   └── wide_resnet_20_350.model
+└── Run5
+    ├── wide_resnet_21_350.model
+    ├── wide_resnet_22_350.model
+    ├── wide_resnet_23_350.model
+    ├── wide_resnet_24_350.model
+    └── wide_resnet_25_350.model
+```
+Following are the main parameters for evaluation:
+```
+--seed: seed used for initializing the first trained model
+--dataset: dataset used for training (cifar10/cifar100)
+--ood_dataset: OoD dataset to compute AUROC
+--load-path: /path/to/pretrained/models/
+--model: model architecture to load (wide_resnet/vgg16)
+--runs: number of experimental runs
+-sn: whether the model was trained using spectral normalization
+--coeff: Coefficient for spectral normalization
+-mod: whether the model was trained using architectural modifications
+--ensemble: number of models in the ensemble
+--model-type: type of model to load for evaluation (softmax/ensemble/gmm)
+```
+As an example, in order to evaluate a Wide-ResNet-28-10 with spectral normalization and architectural modifications on CIFAR-10 with OoD dataset as SVHN, use the following:
+```
+python evaluate.py \
+       --seed 1 \
+       --dataset cifar10 \
+       --ood_dataset svhn \
+       --load-path /path/to/pretrained/models/ \
+       --model wide_resnet \
+       --runs 5 \
+       -sn -mod \
+       --coeff 3.0 \
+       --model-type softmax
+```
+Similarly, to evaluate the above model using feature density, set ```--model-type gmm```. The evaluation script assumes that the seeds of models trained in consecutive runs differ by 1. The script stores the results in a json file with the following structure: 
+```
+{
+    "mean": {
+        "accuracy": mean accuracy,
+        "ece": mean ECE,
+        "m1_auroc": mean AUROC using log density / MI for ensembles,
+        "m1_auprc": mean AUPRC using log density / MI for ensembles,
+        "m2_auroc": mean AUROC using entropy / PE for ensembles,
+        "m2_auprc": mean AUPRC using entropy / PE for ensembles,
+        "t_ece": mean ECE (post temp scaling)
+        "t_m1_auroc": mean AUROC using log density / MI for ensembles (post temp scaling),
+        "t_m1_auprc": mean AUPRC using log density / MI for ensembles (post temp scaling),
+        "t_m2_auroc": mean AUROC using entropy / PE for ensembles (post temp scaling),
+        "t_m2_auprc": mean AUPRC using entropy / PE for ensembles (post temp scaling)
+    },
+    "std": {
+        "accuracy": std error accuracy,
+        "ece": std error ECE,
+        "m1_auroc": std error AUROC using log density / MI for ensembles,
+        "m1_auprc": std error AUPRC using log density / MI for ensembles,
+        "m2_auroc": std error AUROC using entropy / PE for ensembles,
+        "m2_auprc": std error AUPRC using entropy / PE for ensembles,
+        "t_ece": std error ECE (post temp scaling),
+        "t_m1_auroc": std error AUROC using log density / MI for ensembles (post temp scaling),
+        "t_m1_auprc": std error AUPRC using log density / MI for ensembles (post temp scaling),
+        "t_m2_auroc": std error AUROC using entropy / PE for ensembles (post temp scaling),
+        "t_m2_auprc": std error AUPRC using entropy / PE for ensembles (post temp scaling)
+    },
+    "values": {
+        "accuracy": accuracy list,
+        "ece": ece list,
+        "m1_auroc": AUROC list using log density / MI for ensembles,
+        "m2_auroc": AUROC list using entropy / PE for ensembles,
+        "t_ece": ece list (post temp scaling),
+        "t_m1_auroc": AUROC list using log density / MI for ensembles (post temp scaling),
+        "t_m1_auprc": AUPRC list using log density / MI for ensembles (post temp scaling),
+        "t_m2_auroc": AUROC list using entropy / PE for ensembles (post temp scaling),
+        "t_m2_auprc": AUPRC list using entropy / PE for ensembles (post temp scaling)
+    },
+    "info": {dictionary of args}
+}
+```
+
+### Results
+
+#### Dirty-MNIST
+
+To visualise DDU's performance on Dirty-MNIST (i.e., Fig. 1 of the paper), use [fig_1_plot.ipynb](notebooks/fig_1_plot.ipynb). The notebook requires a pretrained LeNet, VGG-16 and ResNet-18 with spectral normalization trained on Dirty-MNIST and visualises the softmax entropy and feature density for Dirty-MNIST (iD) samples vs Fashion-MNIST (OoD) samples. The notebook also visualises the softmax entropies of MNIST vs Ambiguous-MNIST samples for the ResNet-18+SN model (Fig. 2 of the paper). The following figure shows the output of the notebook for the LeNet, VGG-16 and ResNet18+SN model we trained on Dirty-MNIST.
+
+<p align="center">
+  <img src="vis/dirty_mnist_vis.png" width="500" />
+</p>
+
+#### CIFAR-10 vs SVHN
+
+The following table presents results for a Wide-ResNet-28-10 architecture trained on CIFAR-10 with SVHN as the OoD dataset. For the full set of results, refer to the [paper](https://arxiv.org/abs/2102.11582).
+
+| Method  | Aleatoric Uncertainty | Epistemic Uncertainty | Test Accuracy | Test ECE | AUROC |
+| ---  | --- | --- | --- | --- | --- |
+| Softmax  | Softmax Entropy | Softmax Entropy | 95.98+-0.02 | 0.85+-0.02 | 94.44+-0.43 |
+| [Energy-based](https://arxiv.org/abs/2010.03759) | Softmax Entropy | Softmax Density | 95.98+-0.02 | 0.85+-0.02 | 94.56+-0.51 |
+| [5-Ensemble](https://arxiv.org/abs/1612.01474)  | Predictive Entropy | Predictive Entropy | 96.59+-0.02 | 0.76+-0.03 | 97.73+-0.31 |
+| DDU (ours)  | Softmax Entropy | GMM Density | 95.97+-0.03 | 0.85+-0.04 | 98.09+-0.10 |
+
+
+## Active Learning
+
+To run active learning experiments, use ```active_learning_script.py```. You can run active learning experiments on both [MNIST](http://yann.lecun.com/exdb/mnist/) as well as [Dirty-MNIST](https://blackhc.github.io/ddu_dirty_mnist/). When running with Dirty-MNIST, you will need to provide a pretrained model on Dirty-MNIST to distinguish between clean MNIST and Ambiguous-MNIST samples. The following are the main command line arguments for ```active_learning_script.py```.
+```
+--seed: seed used for initializing the first model (later experimental runs will have seeds incremented by 1)
+--model: model architecture to train (resnet18)
+-ambiguous: whether to use ambiguous MNIST during training. If this is set to True, the models will be trained on Dirty-MNIST, otherwise they will train on MNIST.
+--dataset-root: /path/to/amnist_labels.pt and amnist_samples.pt/
+--trained-model: model architecture of pretrained model to distinguish clean and ambiguous MNIST samples
+-tsn: if pretrained model has been trained using spectral normalization
+--tcoeff: coefficient of spectral normalization used on pretrained model
+-tmod: if pretrained model has been trained using architectural modifications (leaky ReLU and average pooling on skip connections)
+--saved-model-path: /path/to/saved/pretrained/model/
+--saved-model-name: name of the saved pretrained model file
+--threshold: Threshold of softmax entropy to decide if a sample is ambiguous (samples having higher softmax entropy than threshold will be considered ambiguous)
+--subsample: number of clean MNIST samples to use to subsample clean MNIST
+-sn: whether to use spectral normalization during training
+--coeff: coefficient of spectral normalization during training
+-mod: whether to use architectural modifications (leaky ReLU and average pooling on skip connections) during training
+--al-type: type of active learning acquisition model (softmax/ensemble/gmm)
+-mi: whether to use mutual information for ensemble al-type
+--num-initial-samples: number of initial samples in the training set
+--max-training-samples: maximum number of training samples
+--acquisition-batch-size: batch size for each acquisition step
+```
+
+As an example, to run the active learning experiment on MNIST using the DDU method, use:
+```
+python active_learning_script.py \
+       --seed 1 \
+       --model resnet18 \
+       -sn -mod \
+       --al-type gmm
+```
+Similarly, to run the active learning experiment on Dirty-MNIST using the DDU baseline, with a pretrained ResNet-18 with SN to distinguish clean and ambiguous MNIST samples, use the following:
+```
+python active_learning_script.py \
+       --seed 1 \
+       --model resnet18 \
+       -sn -mod \
+       -ambiguous \
+       --dataset-root /home/user/amnist/ \
+       --trained-model resnet18 \
+       -tsn \
+       --saved-model-path /path/to/pretrained/model \
+       --saved-model-name resnet18_sn_3.0_1_350.model \
+       --threshold 1.0 \
+       --subsample 1000 \
+       --al-type gmm
+```
+
+### Results
+
+The active learning script stores all results in json files. The MNIST test set accuracy is stored in a json file with the following structure:
+```
+{
+    "experiment run": list of MNIST test set accuracies one per acquisition step
+}
+```
+When using ambiguous samples in the pool set, the script also stores the fraction of ambiguous samples acquired in each step in the following json:
+```
+{
+    "experiment run": list of fractions of ambiguous samples in the acquired training set
+}
+```
+
+### Visualisation
+
+To visualise results from the above json files, use the [al_plot.ipynb](notebooks/al_plot.ipynb) notebook. The following diagram shows the performance of different baselines (softmax, ensemble PE, ensemble MI and DDU) on MNIST and Dirty-MNIST.
+
+<p align="center">
+  <img src="vis/al_plots.png" width="700" />
+</p>
+
+
+## Questions
+
+For any questions, please feel free to raise an issue or email us directly. Our emails can be found on the [paper](https://arxiv.org/abs/2102.11582).

SPC-UQ/Image_Classification/data/ood_detection/cifar10.py

@@ -0,0 +1,107 @@
+"""
+Create train, valid, test iterators for CIFAR-10.
+Train set size: 45000
+Val set size: 5000
+Test set size: 10000
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Subset
+
+from torchvision import datasets
+from torchvision import transforms
+
+
+def get_train_valid_loader(batch_size, augment, val_seed, imagesize=128, val_size=0.1, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Utility function for loading and returning train and valid
+    multi-process iterators over the CIFAR-10 dataset.
+    Params:
+    ------
+    - batch_size: how many samples per batch to load.
+    - augment: whether to apply the data augmentation scheme
+      mentioned in the paper. Only applied on the train split.
+    - val_seed: fix seed for reproducibility.
+    - val_size: percentage split of the training set used for
+      the validation set. Should be a float in the range [0, 1].
+    - num_workers: number of subprocesses to use when loading the dataset.
+    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
+      True if using GPU.
+
+    Returns
+    -------
+    - train_loader: training set iterator.
+    - valid_loader: validation set iterator.
+    """
+    error_msg = "[!] val_size should be in the range [0, 1]."
+    assert (val_size >= 0) and (val_size <= 1), error_msg
+
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transforms
+    valid_transform = transforms.Compose([transforms.Resize(imagesize),transforms.ToTensor(), normalize,])
+
+    if augment:
+        train_transform = transforms.Compose(
+            [transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(),transforms.Resize(imagesize), transforms.ToTensor(), normalize,]
+        )
+    else:
+        train_transform = transforms.Compose([transforms.Resize(imagesize),transforms.ToTensor(), normalize,])
+
+    # load the dataset
+    data_dir = "./data"
+    train_dataset = datasets.CIFAR10(root=data_dir, train=True, download=True, transform=train_transform,)
+
+    valid_dataset = datasets.CIFAR10(root=data_dir, train=True, download=False, transform=valid_transform,)
+
+    num_train = len(train_dataset)
+    indices = list(range(num_train))
+    split = int(np.floor(val_size * num_train))
+
+    np.random.seed(val_seed)
+    np.random.shuffle(indices)
+
+    train_idx, valid_idx = indices[split:], indices[:split]
+
+    train_subset = Subset(train_dataset, train_idx)
+    valid_subset = Subset(valid_dataset, valid_idx)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=True,
+    )
+    valid_loader = torch.utils.data.DataLoader(
+        valid_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=False,
+    )
+
+    return (train_loader, valid_loader)
+
+
+def get_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Utility function for loading and returning a multi-process
+    test iterator over the CIFAR-10 dataset.
+    If using CUDA, num_workers should be set to 1 and pin_memory to True.
+    Params
+    ------
+    - batch_size: how many samples per batch to load.
+    - num_workers: number of subprocesses to use when loading the dataset.
+    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
+      True if using GPU.
+    Returns
+    -------
+    - data_loader: test set iterator.
+    """
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transform
+    transform = transforms.Compose([transforms.Resize(imagesize), transforms.ToTensor(), normalize,])
+
+    data_dir = "./data"
+    dataset = datasets.CIFAR10(root=data_dir, train=False, download=True, transform=transform,)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/cifar100.py

@@ -0,0 +1,107 @@
+"""
+Create train, valid, test iterators for CIFAR-100.
+Train set size: 45000
+Val set size: 5000
+Test set size: 10000
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Subset
+
+from torchvision import datasets
+from torchvision import transforms
+
+
+def get_train_valid_loader(batch_size, augment, val_seed, imagesize=128, val_size=0.1, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Utility function for loading and returning train and valid
+    multi-process iterators over the CIFAR-100 dataset. 
+    Params:
+    ------
+    - batch_size: how many samples per batch to load.
+    - augment: whether to apply the data augmentation scheme
+      mentioned in the paper. Only applied on the train split.
+    - val_seed: fix seed for reproducibility.
+    - val_size: percentage split of the training set used for
+      the validation set. Should be a float in the range [0, 1].
+    - num_workers: number of subprocesses to use when loading the dataset.
+    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
+      True if using GPU.
+
+    Returns
+    -------
+    - train_loader: training set iterator.
+    - valid_loader: validation set iterator.
+    """
+    error_msg = "[!] val_size should be in the range [0, 1]."
+    assert (val_size >= 0) and (val_size <= 1), error_msg
+
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transforms
+    valid_transform = transforms.Compose([transforms.Resize(imagesize),transforms.ToTensor(), normalize,])
+
+    if augment:
+        train_transform = transforms.Compose(
+            [transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.Resize(imagesize),transforms.ToTensor(), normalize,]
+        )
+    else:
+        train_transform = transforms.Compose([transforms.Resize(imagesize),transforms.ToTensor(), normalize,])
+
+    # load the dataset
+    data_dir = "./data"
+    train_dataset = datasets.CIFAR100(root=data_dir, train=True, download=True, transform=train_transform,)
+
+    valid_dataset = datasets.CIFAR100(root=data_dir, train=True, download=False, transform=valid_transform,)
+
+    num_train = len(train_dataset)
+    indices = list(range(num_train))
+    split = int(np.floor(val_size * num_train))
+
+    np.random.seed(val_seed)
+    np.random.shuffle(indices)
+
+    train_idx, valid_idx = indices[split:], indices[:split]
+
+    train_subset = Subset(train_dataset, train_idx)
+    valid_subset = Subset(valid_dataset, valid_idx)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=True,
+    )
+    valid_loader = torch.utils.data.DataLoader(
+        valid_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=False,
+    )
+
+    return (train_loader, valid_loader)
+
+
+def get_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Utility function for loading and returning a multi-process
+    test iterator over the CIFAR-100 dataset.
+    If using CUDA, num_workers should be set to 1 and pin_memory to True.
+    Params
+    ------
+    - batch_size: how many samples per batch to load.
+    - num_workers: number of subprocesses to use when loading the dataset.
+    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
+      True if using GPU.
+    Returns
+    -------
+    - data_loader: test set iterator.
+    """
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transform
+    transform = transforms.Compose([transforms.Resize(imagesize), transforms.ToTensor(), normalize,])
+
+    data_dir = "./data"
+    dataset = datasets.CIFAR100(root=data_dir, train=False, download=True, transform=transform,)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/imagenet.py

@@ -0,0 +1,85 @@
+"""
+Create train, valid, test iterators for ImageNet.
+Train set size: user-defined
+Val set size: user-defined
+Test set size: user-defined (if available)
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Subset
+from torchvision import datasets, transforms
+
+
+def get_train_valid_loader(batch_size, augment, val_seed, imagesize=224, val_size=0.1, num_workers=1, pin_memory=False, **kwargs):
+    assert 0 <= val_size <= 1, "[!] val_size should be in the range [0, 1]."
+
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    imagesize = 224
+    # Define transformations
+    valid_transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(imagesize),
+        transforms.ToTensor(),
+        normalize
+    ])
+    if augment:
+        transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+    else:
+        train_transform = valid_transform
+
+    data_dir = "./data/Imagenet1K"
+    # Load the dataset
+    train_dataset = datasets.ImageFolder(root=f"{data_dir}/train", transform=train_transform)
+    valid_dataset = datasets.ImageFolder(root=f"{data_dir}/train", transform=valid_transform)
+
+    num_train = len(train_dataset)
+    indices = list(range(num_train))
+    split = int(np.floor(val_size * num_train))
+
+    np.random.seed(val_seed)
+    np.random.shuffle(indices)
+
+    train_idx, valid_idx = indices[split:], indices[:split]
+
+    train_subset = Subset(train_dataset, train_idx)
+    valid_subset = Subset(valid_dataset, valid_idx)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=True,
+    )
+    valid_loader = torch.utils.data.DataLoader(
+        valid_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=False,
+    )
+
+    return train_loader, valid_loader
+
+
+def get_test_loader(batch_size, imagesize=224, num_workers=1, pin_memory=False, **kwargs):
+
+    # Define transformation
+    transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+    data_dir = "./data/Imagenet1K"
+    dataset = datasets.ImageFolder(root=f"{data_dir}/val", transform=transform)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/imagenet_a.py

@@ -0,0 +1,37 @@
+"""
+Create train, valid, test iterators for ImageNet.
+Train set size: user-defined
+Val set size: user-defined
+Test set size: user-defined (if available)
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Subset
+from torchvision import datasets, transforms
+
+
+
+def get_test_loader(batch_size, imagesize=224, num_workers=1, pin_memory=False,  **kwargs):
+    transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(imagesize),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+    root = "./data/imagenet-a"
+    dataset = datasets.ImageFolder(
+        root=root,
+        transform=transform
+    )
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False,
+        num_workers=num_workers, pin_memory=pin_memory
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/imagenet_o.py

@@ -0,0 +1,37 @@
+"""
+Create train, valid, test iterators for ImageNet.
+Train set size: user-defined
+Val set size: user-defined
+Test set size: user-defined (if available)
+"""
+
+import torch
+import numpy as np
+from torch.utils.data import Subset
+from torchvision import datasets, transforms
+
+
+
+def get_test_loader(batch_size, imagesize=224, num_workers=1, pin_memory=False,  **kwargs):
+    transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(imagesize),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+    root = "./data/imagenet-o"
+    dataset = datasets.ImageFolder(
+        root=root,
+        transform=transform
+    )
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False,
+        num_workers=num_workers, pin_memory=pin_memory
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/ood_union.py

@@ -0,0 +1,105 @@
+import os
+import torch
+from torch.utils.data import Subset, ConcatDataset, DataLoader
+import random
+
+from torchvision import datasets, transforms
+from torchvision.datasets import ImageFolder
+
+def get_svhn_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transform
+    transform = transforms.Compose([transforms.Resize(imagesize), transforms.ToTensor(), normalize,])
+
+
+    data_dir = "./data"
+    dataset = datasets.SVHN(root=data_dir, split="test", download=True, transform=transform,)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader
+
+
+def get_tinyimagenet_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    transform = transforms.Compose([
+        transforms.Resize(imagesize),
+        transforms.ToTensor(),
+        normalize
+    ])
+
+    data_dir = "./data/tinyimagenet"
+    test_dir = os.path.join(data_dir, "test")
+    dataset = ImageFolder(root=test_dir, transform=transform)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory
+    )
+    return data_loader
+
+
+def get_cifar10_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transform
+    transform = transforms.Compose([transforms.Resize(imagesize), transforms.ToTensor(), normalize,])
+
+    data_dir = "./data"
+    dataset = datasets.CIFAR10(root=data_dir, train=False, download=True, transform=transform,)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader
+
+
+def get_cifar100_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transform
+    transform = transforms.Compose([transforms.Resize(imagesize), transforms.ToTensor(), normalize,])
+
+    data_dir = "./data"
+    dataset = datasets.CIFAR100(root=data_dir, train=False, download=True, transform=transform,)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader
+
+
+
+def get_combined_ood_test_loader(batch_size, sample_seed, imagesize=128, num_workers=1, pin_memory=False, sample_size=10000, **kwargs):
+    svhn_ds = get_svhn_test_loader(batch_size=1, imagesize=imagesize).dataset
+    tiny_ds = get_tinyimagenet_test_loader(batch_size=1, imagesize=imagesize).dataset
+
+    combined_dataset = ConcatDataset([
+        svhn_ds,
+        tiny_ds
+    ])
+
+    # print(len(combined_dataset))
+
+    random.seed(sample_seed)
+    if sample_size is not None and sample_size < len(combined_dataset):
+        indices = random.sample(range(len(combined_dataset)), sample_size)
+        combined_dataset = Subset(combined_dataset, indices)
+
+    data_loader = DataLoader(
+        combined_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/svhn.py

@@ -0,0 +1,94 @@
+import os
+import torch
+import numpy as np
+from torch.utils.data import Subset
+
+from torchvision import datasets
+from torchvision import transforms
+
+
+def get_train_valid_loader(batch_size, augment, val_seed, imagesize=128, val_size=0.1, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Utility function for loading and returning train and valid
+    multi-process iterators over the SVHN dataset. 
+    Params:
+    ------
+    - batch_size: how many samples per batch to load.
+    - augment: whether to apply the data augmentation scheme
+      mentioned in the paper. Only applied on the train split.
+    - val_seed: fix seed for reproducibility.
+    - val_size: percentage split of the training set used for
+      the validation set. Should be a float in the range [0, 1].
+    - num_workers: number of subprocesses to use when loading the dataset.
+    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
+      True if using GPU.
+
+    Returns
+    -------
+    - train_loader: training set iterator.
+    - valid_loader: validation set iterator.
+    """
+    error_msg = "[!] val_size should be in the range [0, 1]."
+    assert (val_size >= 0) and (val_size <= 1), error_msg
+
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transforms
+    valid_transform = transforms.Compose([transforms.Resize(imagesize),transforms.ToTensor(), normalize,])
+
+    # load the dataset
+    data_dir = "./data"
+    train_dataset = datasets.SVHN(root=data_dir, split="train", download=True, transform=valid_transform,)
+
+    valid_dataset = datasets.SVHN(root=data_dir, split="train", download=True, transform=valid_transform,)
+
+    num_train = len(train_dataset)
+    indices = list(range(num_train))
+    split = int(np.floor(val_size * num_train))
+
+    np.random.seed(val_seed)
+    np.random.shuffle(indices)
+
+    train_idx, valid_idx = indices[split:], indices[:split]
+    train_subset = Subset(train_dataset, train_idx)
+    valid_subset = Subset(valid_dataset, valid_idx)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=True,
+    )
+    valid_loader = torch.utils.data.DataLoader(
+        valid_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=False,
+    )
+
+    return (train_loader, valid_loader)
+
+
+def get_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Utility function for loading and returning a multi-process
+    test iterator over the SVHN dataset.
+    If using CUDA, num_workers should be set to 1 and pin_memory to True.
+    Params
+    ------
+    - batch_size: how many samples per batch to load.
+    - num_workers: number of subprocesses to use when loading the dataset.
+    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
+      True if using GPU.
+    Returns
+    -------
+    - data_loader: test set iterator.
+    """
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010],)
+
+    # define transform
+    transform = transforms.Compose([transforms.Resize(imagesize), transforms.ToTensor(), normalize,])
+
+
+    data_dir = "./data"
+    dataset = datasets.SVHN(root=data_dir, split="test", download=True, transform=transform,)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory,
+    )
+
+    return data_loader

SPC-UQ/Image_Classification/data/ood_detection/tinyimagenet.py

@@ -0,0 +1,115 @@
+"""
+Create train, valid, test iterators for Tiny-ImageNet.
+Train set size: 90000 (450 per class)
+Val set size: 10000 (50 per class)
+Test set size: 10000 (no labels)
+"""
+
+import torch
+import numpy as np
+import os
+from torch.utils.data import Subset
+from torchvision import datasets, transforms
+from torchvision.datasets import ImageFolder
+
+
+def get_train_valid_loader(batch_size, augment, val_seed, imagesize=128, val_size=0.1, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Load and return train and valid iterators over the Tiny-ImageNet dataset.
+
+    Params:
+    ------
+    - data_dir: path to Tiny-ImageNet dataset directory.
+    - batch_size: number of samples per batch.
+    - augment: whether to apply data augmentation.
+    - val_seed: random seed for reproducibility.
+    - val_size: fraction of the training set used for validation (0 to 1).
+    - num_workers: number of subprocesses for data loading.
+    - pin_memory: set to True if using GPU.
+
+    Returns:
+    -------
+    - train_loader: training set iterator.
+    - valid_loader: validation set iterator.
+    """
+    assert 0 <= val_size <= 1, "[!] val_size should be in the range [0, 1]."
+
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    # Define transforms
+    valid_transform = transforms.Compose([
+        transforms.Resize(imagesize),
+        transforms.ToTensor(),
+        normalize
+    ])
+
+    if augment:
+        train_transform = transforms.Compose([
+            transforms.RandomCrop(256, padding=4),
+            transforms.RandomHorizontalFlip(),
+            transforms.Resize(imagesize),
+            transforms.ToTensor(),
+            normalize
+        ])
+    else:
+        train_transform = valid_transform
+
+    # Load dataset
+    data_dir = "./data/tinyimagenet"
+    train_dir = os.path.join(data_dir, "train")
+    train_dataset = ImageFolder(root=train_dir, transform=train_transform)
+    valid_dataset = ImageFolder(root=train_dir, transform=valid_transform)  # Same dataset, different transform
+
+    num_train = len(train_dataset)
+    indices = list(range(num_train))
+    split = int(np.floor(val_size * num_train))
+
+    np.random.seed(val_seed)
+    np.random.shuffle(indices)
+
+    train_idx, valid_idx = indices[split:], indices[:split]
+
+    train_subset = Subset(train_dataset, train_idx)
+    valid_subset = Subset(valid_dataset, valid_idx)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=True
+    )
+    valid_loader = torch.utils.data.DataLoader(
+        valid_subset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, shuffle=False
+    )
+
+    return train_loader, valid_loader
+
+
+def get_test_loader(batch_size, imagesize=128, num_workers=1, pin_memory=False, **kwargs):
+    """
+    Load and return a test iterator over the Tiny-ImageNet dataset.
+
+    Params:
+    ------
+    - data_dir: path to Tiny-ImageNet dataset directory.
+    - batch_size: number of samples per batch.
+    - num_workers: number of subprocesses for data loading.
+    - pin_memory: set to True if using GPU.
+
+    Returns:
+    -------
+    - data_loader: test set iterator.
+    """
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    transform = transforms.Compose([
+        transforms.Resize(imagesize),
+        transforms.ToTensor(),
+        normalize
+    ])
+
+    data_dir = "./data/tinyimagenet"
+    test_dir = os.path.join(data_dir, "test")
+    dataset = ImageFolder(root=test_dir, transform=transform)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory
+    )
+    return data_loader

SPC-UQ/Image_Classification/environment.yml

@@ -0,0 +1,16 @@
+name: DDU
+channels:
+    - pytorch
+    - defaults
+dependencies:
+    - python
+    - pytorch=1.7.1
+    - torchvision=0.8.2
+    - cudatoolkit=10.1
+    - tqdm
+    - tensorboard
+    - numpy
+    - scipy
+    - matplotlib
+    - seaborn
+    - scikit-learn

SPC-UQ/Image_Classification/evaluate.py

@@ -0,0 +1,1427 @@
+"""
+Script to evaluate a single model. 
+"""
+import os
+import gc
+import json
+import math
+import torch
+import argparse
+import torch.backends.cudnn as cudnn
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.metrics import accuracy_score
+
+# Import dataloaders
+import data.ood_detection.cifar10 as cifar10
+import data.ood_detection.cifar100 as cifar100
+import data.ood_detection.svhn as svhn
+import data.ood_detection.imagenet as imagenet
+import data.ood_detection.tinyimagenet as tinyimagenet
+import data.ood_detection.imagenet_o as imagenet_o
+import data.ood_detection.imagenet_a as imagenet_a
+import data.ood_detection.ood_union as ood_union
+
+# Import network models
+from net.resnet import resnet50
+from net.resnet_edl import resnet50_edl
+from net.wide_resnet import wrn
+from net.wide_resnet_edl import wrn_edl
+from net.wide_resnet_uq import wrn_uq
+from net.vgg import vgg16
+from net.vgg_edl import vgg16_edl
+from net.vgg_uq import vgg16_uq
+from net.imagenet_wide import imagenet_wide
+from net.imagenet_vgg import imagenet_vgg16
+from net.imagenet_vit import imagenet_vit
+
+# Import metrics to compute
+from metrics.classification_metrics import (
+    test_classification_net,
+    test_classification_net_logits,
+    test_classification_uq,
+    test_classification_net_ensemble,
+    test_classification_net_edl,
+    create_adversarial_dataloader,
+    test_classification_net_logits_edl,
+    test_classification_net_softmax
+)
+from metrics.calibration_metrics import expected_calibration_error
+from metrics.uncertainty_confidence import entropy, logsumexp, self_consistency, edl_unc, certificate
+from metrics.ood_metrics import get_roc_auc, get_roc_auc_logits, get_roc_auc_ensemble, get_unc_ensemble, get_roc_auc_uncs
+from metrics.classification_metrics import get_logits_labels
+from metrics.classification_metrics import get_logits_labels_uq
+
+# Import GMM utils
+from utils.gmm_utils import get_embeddings, gmm_evaluate, gmm_fit
+from utils.ensemble_utils import load_ensemble, Ensemble_fit, Ensemble_evaluate, Ensemble_load
+from utils.oc_utils import oc_fit, oc_evaluate
+from utils.eval_utils import model_load_name
+from utils.train_utils import model_save_name
+from utils.args import eval_args
+
+# Import SPC utils
+from utils.spc_utils import SPC_fit, SPC_load, SPC_evaluate
+
+# Import EDL utils
+from utils.edl_utils import EDL_fit, EDL_load, EDL_evaluate
+
+# Temperature scaling
+from utils.temperature_scaling import ModelWithTemperature
+
+# Dataset params
+dataset_num_classes = {"cifar10": 10, "cifar100": 100, "svhn": 10, "imagenet": 1000, "tinyimagenet": 200, "imagenet_o":200, "imagenet_a":200}
+
+dataset_loader = {"cifar10": cifar10, "cifar100": cifar100, "svhn": svhn, "imagenet": imagenet, "tinyimagenet": tinyimagenet, "imagenet_o":imagenet_o, "imagenet_a":imagenet_a}
+
+# Mapping model name to model function
+models = {"resnet50": resnet50, "resnet50_edl":resnet50_edl, "wide_resnet": wrn, "wide_resnet_edl": wrn_edl, "wide_resnet_uq": wrn_uq, "vgg16": vgg16, "vgg16_edl": vgg16_edl, "vgg16_uq": vgg16_uq, "imagenet_wide":imagenet_wide, "imagenet_vgg16":imagenet_vgg16, "imagenet_vit":imagenet_vit}
+
+model_to_num_dim = {"resnet50": 2048, "resnet50_edl":2048, "wide_resnet": 640, "wide_resnet_edl": 640, "wide_resnet_uq": 640, "vgg16": 512, "vgg16_edl": 512, "vgg16_uq": 512, "imagenet_wide":2048, "imagenet_vgg16":4096, "imagenet_vit":768}
+
+model_to_input_dim = {"resnet50": 32, "resnet50_edl": 32, "wide_resnet": 32, "wide_resnet_edl": 32, "wide_resnet_uq": 32, "vgg16": 32, "vgg16_edl": 32, "vgg16_uq": 32, "imagenet_wide":224, "imagenet_vgg16":224, "imagenet_vit":224}
+
+model_to_last_layer = {"resnet50": "module.fc", "wide_resnet": "module.linear", "vgg16": "module.classifier", "imagenet_wide": "module.linear", "imagenet_vgg16": "module.classifier", "imagenet_vit": "module.linear"}
+
+if __name__ == "__main__":
+
+    args = eval_args().parse_args()
+
+    # Checking if GPU is available
+    cuda = torch.cuda.is_available()
+
+    # Setting additional parameters
+    print("Parsed args", args)
+    print("Seed: ", args.seed)
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda" if cuda else "cpu")
+
+    # Taking input for the dataset
+    num_classes = dataset_num_classes[args.dataset]
+
+    test_loader = dataset_loader[args.dataset].get_test_loader(batch_size=args.batch_size, imagesize=model_to_input_dim[args.model], pin_memory=args.gpu)
+
+    if args.ood_dataset=='ood_union':
+        ood_test_loader = ood_union.get_combined_ood_test_loader(batch_size=args.batch_size, sample_seed=args.seed, imagesize=model_to_input_dim[args.model], pin_memory=args.gpu)
+    else:
+        ood_test_loader = dataset_loader[args.ood_dataset].get_test_loader(batch_size=args.batch_size,
+                                                                           imagesize=model_to_input_dim[args.model],
+                                                                           pin_memory=args.gpu)
+
+    # Evaluating the models
+    accuracies = []
+    c_accuracies = []
+
+    # Pre temperature scaling
+    # m1 - Uncertainty/Confidence Metric 1
+    #      for deterministic model: logsumexp, for ensemble: entropy
+    # m2 - Uncertainty/Confidence Metric 2
+    #      for deterministic model: entropy, for ensemble: MI
+    eces = []
+    ood_m1_aurocs = []
+    ood_m1_auprcs = []
+    ood_m2_aurocs = []
+    ood_m2_auprcs = []
+
+    err_m1_aurocs = []
+    err_m1_auprcs = []
+    err_m2_aurocs = []
+    err_m2_auprcs = []
+
+    adv_ep = 0.02
+    adv_m1_aurocs = []
+    adv_m1_auprcs = []
+    adv_m2_aurocs = []
+    adv_m2_auprcs = []
+
+    # Post temperature scaling
+    t_eces = []
+    t_m1_aurocs = []
+    t_m1_auprcs = []
+    t_m2_aurocs = []
+    t_m2_auprcs = []
+
+    c_eces = []
+
+    adv_unc = np.zeros((args.runs, 9))
+    adv_acc = np.zeros((args.runs, 9))
+
+    topt = None
+
+    for i in range(args.runs):
+        print (f"Evaluating run: {(i+1)}")
+        # Loading the model(s)
+        if args.model_type == "ensemble":
+            if args.dataset == 'imagenet':
+                train_loader, val_loader = dataset_loader[args.dataset].get_train_valid_loader(
+                    batch_size=args.batch_size, imagesize=model_to_input_dim[args.model], augment=args.data_aug,
+                    val_seed=(args.seed + i), val_size=args.val_size, pin_memory=args.gpu)
+                net = models[args.model](pretrained=True, num_classes=1000).cuda()
+                if args.gpu:
+                    net.cuda()
+                    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
+                    cudnn.benchmark = True
+                net.eval()
+
+            else:
+                val_loaders = []
+                for j in range(args.ensemble):
+                    train_loader, val_loader = dataset_loader[args.dataset].get_train_valid_loader(
+                        batch_size=args.batch_size, imagesize=model_to_input_dim[args.model], augment=args.data_aug, val_seed=(args.seed+(5*i)+j), val_size=0.1, pin_memory=args.gpu,
+                    )
+                    val_loaders.append(val_loader)
+                # Evaluate an ensemble
+                ensemble_loc = args.load_loc
+                net_ensemble = load_ensemble(
+                    ensemble_loc=ensemble_loc,
+                    model_name=args.model,
+                    device=device,
+                    num_classes=num_classes,
+                    spectral_normalization=args.sn,
+                    mod=args.mod,
+                    coeff=args.coeff,
+                    seed=(i)
+                )
+
+        else:
+            train_loader, val_loader = dataset_loader[args.dataset].get_train_valid_loader(
+                batch_size=args.batch_size, imagesize=model_to_input_dim[args.model], augment=args.data_aug, val_seed=(args.seed+i), val_size=args.val_size, pin_memory=args.gpu,
+            )
+            if args.dataset == 'imagenet':
+                net = models[args.model](pretrained=True, num_classes=1000).cuda()
+                if args.gpu:
+                    net.cuda()
+                    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
+                    cudnn.benchmark = True
+                net.eval()
+
+            else:
+                if args.val_size==0.1 or (not args.crossval):
+                    saved_model_name = os.path.join(
+                        args.load_loc,
+                        "Run" + str(i + 1),
+                        model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_350.model",
+                    )
+                else:
+                    saved_model_name = os.path.join(
+                        args.load_loc,
+                        "Run" + str(i + 1),
+                        model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_350_0"+str(int(args.val_size*10))+".model",
+                    )
+                print(saved_model_name)
+                net = models[args.model](
+                    spectral_normalization=args.sn, mod=args.mod, coeff=args.coeff, num_classes=num_classes, temp=1.0,
+                )
+                if args.gpu:
+                    net.cuda()
+                    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
+                    cudnn.benchmark = True
+                net.load_state_dict(torch.load(str(saved_model_name)))
+                net.eval()
+
+
+        # Evaluating the UQ method
+        if args.model_type == "ensemble":
+            if args.dataset == 'imagenet':
+                ensemble_model_path = os.path.join(
+                    args.load_loc,
+                    "Run" + str(i + 1),
+                    model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_350_ensemble_model.pth",
+                )
+
+                if os.path.exists(ensemble_model_path):
+                    print(f"Loading existing ensemble_model from {ensemble_model_path}")
+                    Ensemble_model = Ensemble_load(ensemble_model_path, model_to_num_dim[args.model],
+                                                         num_classes, device)
+                else:
+                    if args.model == 'imagenet_vgg16':
+                        embed_path = 'data/imagenet_train_vgg_embedding.pt'
+                        # embed_path = 'data/imagenet_val_vgg_embedding.pt'
+                    if args.model == 'imagenet_wide':
+                        embed_path = 'data/imagenet_train_wide_embedding.pt'
+                        # embed_path = 'data/imagenet_val_wide_embedding.pt'
+                    if args.model == 'imagenet_vit':
+                        embed_path = 'data/imagenet_train_vit_embedding.pt'
+                        # embed_path = 'data/imagenet_val_vit_embedding.pt'
+                    if os.path.exists(embed_path):
+                        data = torch.load(embed_path, map_location=device)
+                        embeddings = data['embeddings']
+                        labels = data['labels']
+                    else:
+                        embeddings, labels = get_embeddings(
+                            net,
+                            train_loader,
+                            num_dim=model_to_num_dim[args.model],
+                            dtype=torch.double,
+                            device=device,
+                            storage_device=device,
+                        )
+                        torch.save({'embeddings': embeddings, 'labels': labels}, embed_path)
+                    Ensemble_model = Ensemble_fit(embeddings, labels, model_to_num_dim[args.model], num_classes, device)
+                    torch.save(Ensemble_model.state_dict(), ensemble_model_path)
+                    print(f"Model saved at {ensemble_model_path}")
+
+                logits, predictive_entropy, mut_info, labels = Ensemble_evaluate(net, Ensemble_model, test_loader, model_to_num_dim[args.model], device)
+                ood_logits, ood_predictive_entropy, ood_mut_info, ood_labels = Ensemble_evaluate(net, Ensemble_model, ood_test_loader, model_to_num_dim[args.model], device)
+                (conf_matrix, accuracy, labels_list, predictions, confidences,) = test_classification_net_softmax(logits, labels)
+                t_accuracy = accuracy
+                ece = expected_calibration_error(confidences, predictions, labels_list, num_bins=15)
+                t_ece = ece
+                (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_uncs(predictive_entropy, ood_predictive_entropy, device)
+                (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc_uncs(mut_info, ood_mut_info, device)
+
+                labels_array = np.array(labels_list)
+                pred_array = np.array(predictions)
+                correct_mask = labels_array == pred_array
+                entropy_right = predictive_entropy[correct_mask]
+                entropy_wrong = predictive_entropy[~correct_mask]
+                mut_info_right = mut_info[correct_mask]
+                mut_info_wrong = mut_info[~correct_mask]
+                (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc =  get_roc_auc_uncs(entropy_right, entropy_wrong, device)
+                (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc =  get_roc_auc_uncs(mut_info_right, mut_info_wrong, device)
+
+                adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                               imagesize=model_to_input_dim[args.model],
+                                                                               pin_memory=args.gpu)
+
+                adv_logits, adv_predictive_entropy, adv_mut_info, adv_labels = Ensemble_evaluate(net, Ensemble_model, adv_test_loader, model_to_num_dim[args.model], device)
+
+
+                (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_uncs(predictive_entropy, adv_predictive_entropy, device)
+                (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_uncs(mut_info, adv_mut_info, device)
+
+                print('adv_m1_auroc', adv_m1_auroc)
+
+                t_m1_auroc = ood_m1_auroc
+                t_m1_auprc = ood_m1_auprc
+                t_m2_auroc = ood_m2_auroc
+                t_m2_auprc = ood_m2_auprc
+
+
+            else:
+                (conf_matrix, accuracy, labels_list, predictions, confidences,) = test_classification_net_ensemble(
+                    net_ensemble, test_loader, device
+                )
+                ece = expected_calibration_error(confidences, predictions, labels_list, num_bins=15)
+
+                (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_ensemble(
+                    net_ensemble, test_loader, ood_test_loader, "entropy", device
+                )
+                (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc_ensemble(
+                    net_ensemble, test_loader, ood_test_loader, "mutual_information", device
+                )
+
+                labels_array = np.array(labels_list)
+                pred_array = np.array(predictions)
+                correct_mask = labels_array == pred_array
+                from torch.utils.data import Subset, DataLoader
+                dataset = test_loader.dataset
+                correct_indices = np.where(correct_mask)[0]
+                right_subset = Subset(dataset, correct_indices)
+                right_loader = DataLoader(right_subset, batch_size=test_loader.batch_size, shuffle=False)
+                wrong_indices = np.where(~correct_mask)[0]
+                wrong_subset = Subset(dataset, wrong_indices)
+                wrong_loader = DataLoader(wrong_subset, batch_size=test_loader.batch_size, shuffle=False)
+                (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_ensemble(
+                    net_ensemble, right_loader, wrong_loader, "entropy", device
+                )
+                (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc = get_roc_auc_ensemble(
+                    net_ensemble, right_loader, wrong_loader, "mutual_information", device
+                )
+
+                adv_test_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,batch_size=args.batch_size)
+                (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_ensemble(
+                    net_ensemble, test_loader, adv_test_loader, "entropy", device
+                )
+                (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_ensemble(
+                    net_ensemble, test_loader, adv_test_loader, "mutual_information", device
+                )
+                print('adv_m1_auroc,adv_m2_auroc', adv_m1_auroc, adv_m2_auroc)
+
+                if args.sample_noise:
+                    adv_eps = np.linspace(0, 0.4, 9)
+                    print(adv_eps)
+                    for idx_ep, ep in enumerate(adv_eps):
+                        adv_loader = create_adversarial_dataloader(net_ensemble[0], test_loader, device, epsilon=ep,
+                                                                   batch_size=args.batch_size)
+                        (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions, adv_confidences) = test_classification_net_ensemble(
+                            net_ensemble, adv_loader, device
+                        )
+                        uncertainties = get_unc_ensemble(net_ensemble, adv_loader, "entropy", device).detach().cpu().numpy()
+                        quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                        quantiles = np.delete(quantiles, 0)
+                        unc_list = []
+                        accuracy_list = []
+                        for threshold in quantiles:
+                            cer_indices = (uncertainties < threshold)
+                            unc_indices = ~cer_indices
+                            labels_list = np.array(adv_labels_list)
+                            targets_cer = labels_list[cer_indices]
+                            predictions = np.array(adv_predictions)
+                            pred_cer = predictions[cer_indices]
+                            targets_unc = labels_list[unc_indices]
+                            pred_unc = predictions[unc_indices]
+                            cer_right = np.sum(targets_cer == pred_cer)
+                            cer = len(targets_cer)
+                            unc_right = np.sum(targets_unc == pred_unc)
+                            unc = len(targets_unc)
+                            accuracy_cer = cer_right / cer
+                            accuracy_unc = unc_right / unc
+                            unc_list.append(threshold)
+                            accuracy_list.append(accuracy_cer)
+                            print('ACC:', accuracy_cer, accuracy_unc)
+                        from scipy.stats import spearmanr
+
+                        Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                        print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                        adv_unc[i][idx_ep] = uncertainties.mean()
+                        adv_acc[i][idx_ep] = adv_accuracy
+
+                        (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_ensemble(
+                            net_ensemble, test_loader, adv_loader, "entropy", device
+                        )
+                        (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_ensemble(
+                            net_ensemble, test_loader, adv_loader, "mutual_information", device
+                        )
+                        print('adv_m1_auroc,adv_m2_auroc', adv_m1_auroc, adv_m2_auroc)
+
+                # Temperature scale the ensemble
+                t_ensemble = []
+                for model, val_loader in zip(net_ensemble, val_loaders):
+                    t_model = ModelWithTemperature(model)
+                    t_model.set_temperature(val_loader)
+                    t_ensemble.append(t_model)
+
+                (
+                    t_conf_matrix,
+                    t_accuracy,
+                    t_labels_list,
+                    t_predictions,
+                    t_confidences,
+                ) = test_classification_net_ensemble(t_ensemble, test_loader, device)
+                t_ece = expected_calibration_error(t_confidences, t_predictions, t_labels_list, num_bins=15)
+
+                (_, _, _), (_, _, _), t_m1_auroc, t_m1_auprc = get_roc_auc_ensemble(
+                    t_ensemble, test_loader, ood_test_loader, "entropy", device
+                )
+                (_, _, _), (_, _, _), t_m2_auroc, t_m2_auprc = get_roc_auc_ensemble(
+                    t_ensemble, test_loader, ood_test_loader, "mutual_information", device
+                )
+
+        elif args.model_type == "edl":
+            if args.dataset == 'imagenet':
+                edl_model_path = os.path.join(
+                    args.load_loc,
+                    "Run" + str(i + 1),
+                    model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_350_edl_model.pth",
+                )
+
+                if os.path.exists(edl_model_path):
+                    print(f"Loading existing edl_model from {edl_model_path}")
+                    EDL_model = EDL_load(edl_model_path, model_to_num_dim[args.model],
+                                                         num_classes, device)
+                else:
+                    if args.model=='imagenet_vgg16':
+                        embed_path = 'data/imagenet_train_vgg_embedding.pt'
+                        # embed_path = 'data/imagenet_val_vgg_embedding.pt'
+                    if args.model=='imagenet_wide':
+                        embed_path = 'data/imagenet_train_wide_embedding.pt'
+                        # embed_path = 'data/imagenet_val_wide_embedding.pt'
+                    if args.model=='imagenet_vit':
+                        embed_path = 'data/imagenet_train_vit_embedding.pt'
+                        # embed_path = 'data/imagenet_val_vit_embedding.pt'
+                    if os.path.exists(embed_path):
+                        data = torch.load(embed_path, map_location=device)
+                        embeddings = data['embeddings']
+                        labels = data['labels']
+                    else:
+                        embeddings, labels = get_embeddings(
+                            net,
+                            train_loader,
+                            num_dim=model_to_num_dim[args.model],
+                            dtype=torch.double,
+                            device=device,
+                            storage_device=device,
+                        )
+                        torch.save({'embeddings': embeddings, 'labels': labels}, embed_path)
+                    EDL_model = EDL_fit(embeddings, labels, model_to_num_dim[args.model], num_classes,
+                                                        device)
+                    torch.save(EDL_model.state_dict(), edl_model_path)
+                    print(f"Model saved at {edl_model_path}")
+
+                logits, labels = EDL_evaluate(net, EDL_model, test_loader, model_to_num_dim[args.model], device)
+                ood_logits, ood_labels = EDL_evaluate(net, EDL_model, ood_test_loader, model_to_num_dim[args.model], device)
+                (conf_matrix, accuracy, labels_list, predictions, confidences,) = test_classification_net_logits_edl(logits, labels)
+                t_accuracy = accuracy
+
+                ece = expected_calibration_error(confidences, predictions, labels_list, num_bins=15)
+                t_ece = ece
+
+                (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_logits(logits, ood_logits, edl_unc, device)
+
+                labels_array = np.array(labels_list)
+                pred_array = np.array(predictions)
+                correct_mask = labels_array == pred_array
+                # logits, _ = get_logits_labels(net, test_loader, device)
+                logits_right = logits[correct_mask]
+                logits_wrong = logits[~correct_mask]
+                (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(logits_right, logits_wrong,
+                                                                                      edl_unc, device)
+
+                adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                                   imagesize=model_to_input_dim[args.model],
+                                                                                   pin_memory=args.gpu)
+                # (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,adv_confidences,) = test_classification_net_edl(net, adv_test_loader, device)
+                # adv_logits, _ = get_logits_labels(net, adv_test_loader, device)
+                adv_logits, adv_labels = EDL_evaluate(net, EDL_model, adv_test_loader, model_to_num_dim[args.model], device)
+
+                (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(logits, adv_logits, edl_unc, device)
+                print('adv_m1_auroc', adv_m1_auroc)
+
+                if args.sample_noise:
+                    adv_eps = np.linspace(0, 0.4, 9)
+                    for idx_ep, ep in enumerate(adv_eps):
+                        adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=ep,
+                                                                   batch_size=args.batch_size)
+                        (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                         adv_confidences,) = test_classification_net(net, adv_loader, device)
+                        adv_logits, adv_labels = EDL_evaluate(net, EDL_model, adv_loader,
+                                                                    model_to_num_dim[args.model], device)
+                        uncertainties = edl_unc(adv_logits).detach().cpu().numpy()
+                        quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                        quantiles = np.delete(quantiles, 0)
+                        unc_list = []
+                        accuracy_list = []
+                        for threshold in quantiles:
+                            cer_indices = (uncertainties < threshold)
+                            unc_indices = ~cer_indices
+                            labels_list = np.array(adv_labels_list)
+                            targets_cer = labels_list[cer_indices]
+                            predictions = np.array(adv_predictions)
+                            pred_cer = predictions[cer_indices]
+                            targets_unc = labels_list[unc_indices]
+                            pred_unc = predictions[unc_indices]
+                            cer_right = np.sum(targets_cer == pred_cer)
+                            cer = len(targets_cer)
+                            unc_right = np.sum(targets_unc == pred_unc)
+                            unc = len(targets_unc)
+                            accuracy_cer = cer_right / cer
+                            accuracy_unc = unc_right / unc
+                            unc_list.append(threshold)
+                            accuracy_list.append(accuracy_cer)
+                            print('ACC:', accuracy_cer, accuracy_unc)
+                        from scipy.stats import spearmanr
+
+                        Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                        print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                        adv_unc[i][idx_ep] = uncertainties.mean()
+                        adv_acc[i][idx_ep] = adv_accuracy
+
+
+            else:
+                (conf_matrix, accuracy, labels_list, predictions, confidences,) = test_classification_net_edl(
+                    net, test_loader, device
+                )
+                t_accuracy = accuracy
+                ece = expected_calibration_error(confidences, predictions, labels_list, num_bins=15)
+                t_ece=ece
+
+                (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc(net, test_loader, ood_test_loader, edl_unc, device)
+
+                labels_array = np.array(labels_list)
+                pred_array = np.array(predictions)
+                correct_mask = labels_array == pred_array
+                logits, _ = get_logits_labels(net, test_loader, device)
+                logits_right = logits[correct_mask]
+                logits_wrong = logits[~correct_mask]
+                (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(logits_right, logits_wrong, edl_unc, device)
+
+                adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,
+                                                           batch_size=args.batch_size, edl=True)
+                (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                 adv_confidences,) = test_classification_net_edl(net, adv_loader, device)
+                adv_logits, _ = get_logits_labels(net, adv_loader, device)
+
+                (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(logits, adv_logits, edl_unc, device)
+                print('adv_m1_auroc', adv_m1_auroc)
+
+
+                if args.sample_noise:
+                    adv_eps = np.linspace(0, 0.4, 9)
+                    for idx_ep, ep in enumerate(adv_eps):
+                        adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=ep,
+                                                                   batch_size=args.batch_size, edl=True)
+                        (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                         adv_confidences,) = test_classification_net_edl(net, adv_loader, device)
+                        adv_logits, _ = get_logits_labels(net, adv_loader, device)
+                        uncertainties = edl_unc(adv_logits).detach().cpu().numpy()
+                        quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                        quantiles = np.delete(quantiles, 0)
+                        unc_list = []
+                        accuracy_list = []
+                        for threshold in quantiles:
+                            cer_indices = (uncertainties < threshold)
+                            unc_indices = ~cer_indices
+                            labels_list = np.array(adv_labels_list)
+                            targets_cer = labels_list[cer_indices]
+                            predictions = np.array(adv_predictions)
+                            pred_cer = predictions[cer_indices]
+                            targets_unc = labels_list[unc_indices]
+                            pred_unc = predictions[unc_indices]
+                            cer_right = np.sum(targets_cer == pred_cer)
+                            cer = len(targets_cer)
+                            unc_right = np.sum(targets_unc == pred_unc)
+                            unc = len(targets_unc)
+                            accuracy_cer = cer_right / cer
+                            accuracy_unc = unc_right / unc
+                            unc_list.append(threshold)
+                            accuracy_list.append(accuracy_cer)
+                            print('ACC:', accuracy_cer, accuracy_unc)
+                        from scipy.stats import spearmanr
+
+                        Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                        print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                        adv_unc[i][idx_ep] = uncertainties.mean()
+                        adv_acc[i][idx_ep] = adv_accuracy
+
+            ood_m2_auroc=ood_m1_auroc
+            ood_m2_auprc = ood_m1_auprc
+            err_m2_auroc = err_m1_auroc
+            err_m2_auprc = err_m1_auprc
+            adv_m2_auroc = adv_m1_auroc
+            adv_m2_auprc = adv_m1_auprc
+            t_m1_auroc=ood_m1_auroc
+            t_m1_auprc=ood_m1_auprc
+            t_m2_auroc=ood_m1_auroc
+            t_m2_auprc=ood_m1_auprc
+
+        elif args.model_type == "joint":
+            (conf_matrix, accuracy, labels_list, predictions, confidences,) = test_classification_uq(
+                net, test_loader, device
+            )
+            ece = expected_calibration_error(confidences, predictions, labels_list, num_bins=15)
+            print(accuracy)
+            print('ece', ece)
+
+            t_ece=ece
+            t_accuracy=accuracy
+
+            print("SPC Model")
+
+            logits, labels = get_logits_labels_uq(net, test_loader, device)
+
+            soft = torch.nn.functional.softmax(logits[0], dim=1)
+            delta = torch.min(torch.min(logits[2] - logits[3], logits[1] - 2 * logits[3]), 2 * logits[2] - logits[1])
+
+            uncertainty = abs(logits[2] + logits[3] - logits[1])
+
+            threshold = 0.05
+            mask = (uncertainty < threshold).float()
+            delta = delta * mask
+            softmax_prob = soft + delta
+
+            c_confidences, c_predictions = torch.max(softmax_prob, dim=1)
+            c_predictions = c_predictions.tolist()
+            c_confidences = c_confidences.tolist()
+            c_ece = expected_calibration_error(c_confidences, c_predictions, labels_list, num_bins=15)
+            print('ece', ece, 't_ece', t_ece, 'c_ece', c_ece)
+            c_accuracy = accuracy_score(labels_list, c_predictions)
+            print(accuracy, t_accuracy, c_accuracy)
+
+            ood_logits, ood_labels = get_logits_labels_uq(net, ood_test_loader, device)
+
+            (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_logits(logits, ood_logits, self_consistency,
+                                                                                  device)
+            (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc_logits(logits[0], ood_logits[0], entropy, device)
+
+            labels_array = np.array(labels_list)
+            pred_array = np.array(predictions)
+            correct_mask = labels_array == pred_array
+            logits_right = [m[correct_mask] for m in logits]
+            logits_wrong = [m[~correct_mask] for m in logits]
+            (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(logits_right, logits_wrong,
+                                                                                  self_consistency, device)
+            (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc = get_roc_auc_logits(logits_right[0], logits_wrong[0], entropy,
+                                                                                  device)
+
+            if args.dataset == 'imagenet':
+                adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                               imagesize=model_to_input_dim[args.model],
+                                                                               pin_memory=args.gpu)
+            else:
+                adv_test_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,
+                                                                batch_size=args.batch_size, joint=True)
+
+            adv_logits, adv_labels = get_logits_labels_uq(net, adv_test_loader, device)
+
+            (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(logits, adv_logits, self_consistency, device)
+            (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits[0], adv_logits[0], entropy, device)
+
+            t_m1_auroc = ood_m1_auroc
+            t_m1_auprc = ood_m1_auprc
+            t_m2_auroc = ood_m2_auroc
+            t_m2_auprc = ood_m2_auprc
+
+        else:
+            (conf_matrix, accuracy, labels_list, predictions, confidences,) = test_classification_net(
+                net, test_loader, device
+            )
+            ece = expected_calibration_error(confidences, predictions, labels_list, num_bins=15)
+            print(accuracy)
+            print('ece',ece)
+
+            temp_scaled_net = ModelWithTemperature(net)
+            temp_scaled_net.set_temperature(val_loader)
+            # temp_scaled_net.set_temperature(train_loader)
+            topt = temp_scaled_net.temperature
+
+            (t_conf_matrix, t_accuracy, t_labels_list, t_predictions, t_confidences,) = test_classification_net(
+                temp_scaled_net, test_loader, device
+            )
+            t_ece = expected_calibration_error(t_confidences, t_predictions, t_labels_list, num_bins=15)
+            print('t_ece',t_ece)
+
+            if (args.model_type == "gmm"):
+                # Evaluate a GMM model
+                print("GMM Model")
+
+                if args.crossval:
+                    embeddings, labels = get_embeddings(
+                        net,
+                        val_loader,
+                        num_dim=model_to_num_dim[args.model],
+                        dtype=torch.double,
+                        device=device,
+                        storage_device=device,
+                    )
+                else:
+                    if args.dataset == 'imagenet':
+                        if args.model == 'imagenet_vgg16':
+                            embed_path = 'data/imagenet_train_vgg_embedding.pt'
+                            # embed_path = 'data/imagenet_val_vgg_embedding.pt'
+                        if args.model == 'imagenet_wide':
+                            embed_path = 'data/imagenet_train_wide_embedding.pt'
+                            # embed_path = 'data/imagenet_val_wide_embedding.pt'
+                        if args.model == 'imagenet_vit':
+                            embed_path = 'data/imagenet_train_vit_embedding.pt'
+                            # embed_path = 'data/imagenet_val_vit_embedding.pt'
+                        if os.path.exists(embed_path):
+                            data = torch.load(embed_path, map_location=device)
+                            embeddings = data['embeddings']
+                            labels = data['labels']
+                        else:
+                            embeddings, labels = get_embeddings(
+                                net,
+                                train_loader,
+                                num_dim=model_to_num_dim[args.model],
+                                dtype=torch.double,
+                                device=device,
+                                storage_device=device,
+                            )
+                            torch.save({'embeddings': embeddings, 'labels': labels}, embed_path)
+                    else:
+                        embeddings, labels = get_embeddings(
+                            net,
+                            train_loader,
+                            num_dim=model_to_num_dim[args.model],
+                            dtype=torch.double,
+                            device=device,
+                            storage_device=device,
+                        )
+
+                try:
+                    gaussians_model, jitter_eps = gmm_fit(embeddings=embeddings, labels=labels, num_classes=num_classes, device=device)
+                    logits, labels = gmm_evaluate(
+                        net, gaussians_model, test_loader, device=device, num_classes=num_classes, storage_device=device,
+                    )
+
+                    ood_logits, ood_labels = gmm_evaluate(
+                        net, gaussians_model, ood_test_loader, device=device, num_classes=num_classes, storage_device=device,
+                    )
+
+                    (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_logits(
+                        logits, ood_logits, logsumexp, device, confidence=True
+                    )
+                    (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc_logits(logits, ood_logits, entropy, device)
+
+                    labels_array = np.array(labels_list)
+                    pred_array = np.array(predictions)
+                    correct_mask = labels_array == pred_array
+                    logits_right = logits[correct_mask]
+                    logits_wrong = logits[~correct_mask]
+                    (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(logits_right, logits_wrong,logsumexp, device, confidence=True)
+                    (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc = get_roc_auc_logits(logits_right, logits_wrong,entropy, device, confidence=True)
+
+                    if args.dataset == 'imagenet':
+                        adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                                       imagesize=model_to_input_dim[args.model],
+                                                                                       pin_memory=args.gpu)
+                    else:
+                        adv_test_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,batch_size=args.batch_size)
+                    (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions, adv_confidences,) = test_classification_net(net, adv_test_loader, device)
+                    adv_logits, adv_labels = gmm_evaluate(net, gaussians_model, adv_test_loader, device=device, num_classes=num_classes, storage_device=device, )
+                    (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(logits, adv_logits, logsumexp, device, confidence=True)
+                    (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits, adv_logits, entropy, device)
+
+                    if args.sample_noise:
+                        adv_eps = np.linspace(0, 0.4, 9)
+                        for idx_ep, ep in enumerate(adv_eps):
+                            adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=ep,
+                                                                       batch_size=args.batch_size)
+                            (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                             adv_confidences,) = test_classification_net(net, adv_loader, device)
+                            adv_logits, adv_labels = gmm_evaluate(net, gaussians_model, adv_loader, device=device, num_classes=num_classes,storage_device=device,)
+                            uncertainties = -logsumexp(adv_logits).detach().cpu().numpy()
+                            quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                            quantiles = np.delete(quantiles, 0)
+                            unc_list = []
+                            accuracy_list = []
+                            for threshold in quantiles:
+                                cer_indices = (uncertainties < threshold)
+                                unc_indices = ~cer_indices
+                                labels_list = np.array(adv_labels_list)
+                                targets_cer = labels_list[cer_indices]
+                                predictions = np.array(adv_predictions)
+                                pred_cer = predictions[cer_indices]
+                                targets_unc = labels_list[unc_indices]
+                                pred_unc = predictions[unc_indices]
+                                cer_right = np.sum(targets_cer == pred_cer)
+                                cer = len(targets_cer)
+                                unc_right = np.sum(targets_unc == pred_unc)
+                                unc = len(targets_unc)
+                                accuracy_cer = cer_right / cer
+                                accuracy_unc = unc_right / unc
+                                unc_list.append(threshold)
+                                accuracy_list.append(accuracy_cer)
+                                print('ACC:', accuracy_cer, accuracy_unc)
+                            from scipy.stats import spearmanr
+                            Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                            print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                            adv_unc[i][idx_ep]=uncertainties.mean()
+                            adv_acc[i][idx_ep] = adv_accuracy
+
+                            (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(logits, adv_logits, logsumexp, device, confidence=True)
+                            (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits, adv_logits, entropy, device)
+                            print('adv_m1_auroc,adv_m2_auroc', adv_m1_auroc, adv_m2_auroc)
+
+                    t_m1_auroc = ood_m1_auroc
+                    t_m1_auprc = ood_m1_auprc
+                    t_m2_auroc = ood_m2_auroc
+                    t_m2_auprc = ood_m2_auprc
+
+                except RuntimeError as e:
+                    print("Runtime Error caught: " + str(e))
+                    continue
+
+            elif (args.model_type == "oc"):
+                # Evaluate a OC model
+                print("OC Model")
+
+                if args.crossval:
+                    embeddings, labels = get_embeddings(
+                        net,
+                        val_loader,
+                        num_dim=model_to_num_dim[args.model],
+                        dtype=torch.double,
+                        device=device,
+                        storage_device=device,
+                    )
+                else:
+                    if args.dataset == 'imagenet':
+                        if args.model == 'imagenet_vgg16':
+                            embed_path = 'data/imagenet_train_vgg_embedding.pt'
+                            # embed_path = 'data/imagenet_val_vgg_embedding.pt'
+                        if args.model == 'imagenet_wide':
+                            embed_path = 'data/imagenet_train_wide_embedding.pt'
+                            # embed_path = 'data/imagenet_val_wide_embedding.pt'
+                        if args.model == 'imagenet_vit':
+                            embed_path = 'data/imagenet_train_vit_embedding.pt'
+                            # embed_path = 'data/imagenet_val_vit_embedding.pt'
+                        if os.path.exists(embed_path):
+                            data = torch.load(embed_path, map_location=device)
+                            embeddings = data['embeddings']
+                            labels = data['labels']
+                        else:
+                            embeddings, labels = get_embeddings(
+                                net,
+                                train_loader,
+                                num_dim=model_to_num_dim[args.model],
+                                dtype=torch.double,
+                                device=device,
+                                storage_device=device,
+                            )
+                            torch.save({'embeddings': embeddings, 'labels': labels}, embed_path)
+                    else:
+                        embeddings, labels = get_embeddings(
+                            net,
+                            train_loader,
+                            num_dim=model_to_num_dim[args.model],
+                            dtype=torch.double,
+                            device=device,
+                            storage_device=device,
+                        )
+
+                try:
+                    oc_model = oc_fit(embeddings=embeddings, device=device)
+                    logits, OCs = oc_evaluate(
+                        net, oc_model, test_loader,model_to_num_dim[args.model], device=device
+                    )
+
+                    ood_logits, ood_OCs = oc_evaluate(
+                        net, oc_model, ood_test_loader, model_to_num_dim[args.model], device=device)
+
+                    (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_logits(OCs, ood_OCs, certificate, device, confidence=True)
+                    (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc_logits(logits, ood_logits, entropy, device)
+
+                    labels_array = np.array(labels_list)
+                    pred_array = np.array(predictions)
+                    correct_mask = labels_array == pred_array
+                    logits_right = logits[correct_mask]
+                    logits_wrong = logits[~correct_mask]
+                    OCs_right = OCs[correct_mask]
+                    OCs_wrong = OCs[~correct_mask]
+                    (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(OCs_right, OCs_wrong, certificate, device, confidence=True)
+                    (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc = get_roc_auc_logits(logits_right, logits_wrong, entropy, device)
+
+                    if args.dataset == 'imagenet':
+                        adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                                       imagesize=model_to_input_dim[args.model],
+                                                                                       pin_memory=args.gpu)
+                    else:
+                        adv_test_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,batch_size=args.batch_size)
+                    (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions, adv_confidences,) = test_classification_net(net, adv_test_loader, device)
+                    adv_logits, adv_OCs = oc_evaluate(net, oc_model, adv_test_loader, model_to_num_dim[args.model], device=device)
+                    (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(OCs, adv_OCs, certificate, device, confidence=True)
+                    (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits, adv_logits, entropy, device)
+
+                    if args.sample_noise:
+                        adv_eps = np.linspace(0, 0.4, 9)
+                        for idx_ep, ep in enumerate(adv_eps):
+                            adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=ep,
+                                                                       batch_size=args.batch_size)
+                            (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                             adv_confidences,) = test_classification_net(net, adv_loader, device)
+                            adv_logits, adv_OCs = oc_evaluate(net, oc_model, adv_loader, model_to_num_dim[args.model], device=device)
+                            uncertainties = -adv_OCs.cpu().numpy()
+                            quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                            quantiles = np.delete(quantiles, 0)
+                            unc_list = []
+                            accuracy_list = []
+                            for threshold in quantiles:
+                                cer_indices = (uncertainties < threshold)
+                                unc_indices = ~cer_indices
+                                labels_list = np.array(adv_labels_list)
+                                targets_cer = labels_list[cer_indices]
+                                predictions = np.array(adv_predictions)
+                                pred_cer = predictions[cer_indices]
+                                targets_unc = labels_list[unc_indices]
+                                pred_unc = predictions[unc_indices]
+                                cer_right = np.sum(targets_cer == pred_cer)
+                                cer = len(targets_cer)
+                                unc_right = np.sum(targets_unc == pred_unc)
+                                unc = len(targets_unc)
+                                accuracy_cer = cer_right / cer
+                                accuracy_unc = unc_right / unc
+                                unc_list.append(threshold)
+                                accuracy_list.append(accuracy_cer)
+                                print('ACC:', accuracy_cer, accuracy_unc)
+                            from scipy.stats import spearmanr
+                            Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                            print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                            adv_unc[i][idx_ep]=uncertainties.mean()
+                            adv_acc[i][idx_ep] = adv_accuracy
+
+                            (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(OCs, adv_OCs, certificate, device, confidence=True)
+                            (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits, adv_logits, entropy, device)
+                            print('adv_m1_auroc,adv_m2_auroc', adv_m1_auroc, adv_m2_auroc)
+
+                    t_m1_auroc = ood_m1_auroc
+                    t_m1_auprc = ood_m1_auprc
+                    t_m2_auroc = ood_m2_auroc
+                    t_m2_auprc = ood_m2_auprc
+
+                except RuntimeError as e:
+                    print("Runtime Error caught: " + str(e))
+                    continue
+
+            elif (args.model_type == "spc"):
+                print("SPC Model")
+
+                if args.crossval:
+                    spc_model_path = os.path.join(
+                        args.load_loc,
+                        "Run" + str(i + 1),
+                        model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_valsize_"+str(args.val_size)+"_350_mar_model.pth",
+                    )
+                else:
+                    spc_model_path = os.path.join(
+                        args.load_loc,
+                        "Run" + str(i + 1),
+                        model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_350_mar_model.pth",
+                    )
+
+                if os.path.exists(spc_model_path):
+                    print(f"Loading existing spc_model from {spc_model_path}")
+                    SPC_model = SPC_load(spc_model_path, model_to_num_dim[args.model], num_classes, device)
+                else:
+                    print(f"Model not found. Training a new one...")
+                    if args.crossval:
+                        embeddings, labels = get_embeddings(
+                            net,
+                            val_loader,
+                            num_dim=model_to_num_dim[args.model],
+                            dtype=torch.double,
+                            device=device,
+                            storage_device=device,
+                        )
+                    else:
+                        if args.dataset == 'imagenet':
+                            if args.model=='imagenet_vgg16':
+                                embed_path = 'data/imagenet_train_vgg_embedding.pt'
+                                # embed_path = 'data/imagenet_val_vgg_embedding.pt'
+                            if args.model=='imagenet_wide':
+                                embed_path = 'data/imagenet_train_wide_embedding.pt'
+                                # embed_path = 'data/imagenet_val_wide_embedding.pt'
+                            if args.model=='imagenet_vit':
+                                embed_path = 'data/imagenet_train_vit_embedding.pt'
+                                # embed_path = 'data/imagenet_val_vit_embedding.pt'
+                            if os.path.exists(embed_path):
+                                data = torch.load(embed_path, map_location=device)
+                                embeddings = data['embeddings']
+                                labels = data['labels']
+                            else:
+                                embeddings, labels = get_embeddings(
+                                    net,
+                                    train_loader,
+                                    num_dim=model_to_num_dim[args.model],
+                                    dtype=torch.double,
+                                    device=device,
+                                    storage_device=device,
+                                )
+                                torch.save({'embeddings': embeddings, 'labels': labels}, embed_path)
+                        else:
+                            embeddings, labels = get_embeddings(
+                                net,
+                                train_loader,
+                                num_dim=model_to_num_dim[args.model],
+                                dtype=torch.double,
+                                device=device,
+                                storage_device=device,
+                            )
+
+                    parts= model_to_last_layer[args.model].split('.')
+                    net_last_layer = net
+
+                    for attr in parts:
+                        net_last_layer = getattr(net_last_layer, attr)
+
+                    SPC_model=SPC_fit(net_last_layer, topt, embeddings, labels, model_to_num_dim[args.model], num_classes, device)
+                    torch.save(SPC_model.state_dict(), spc_model_path)
+                    print(f"Model saved at {spc_model_path}")
+
+                logits,mars=SPC_evaluate(net, SPC_model, test_loader, model_to_num_dim[args.model], num_classes, device)
+
+                soft = torch.nn.functional.softmax(logits, dim=1)
+                delta = torch.min(torch.min(mars[2] - mars[3], mars[1] - 2 * mars[3]), 2 * mars[2] - mars[1])
+                # delta=mars[2] - mars[3]
+
+                uncertainty = abs(mars[2]+mars[3] - mars[1])
+
+                threshold=0.05
+                mask=(uncertainty<threshold).float()
+                delta = delta*mask
+                softmax_prob = soft + delta
+                print(torch.sum(softmax_prob, dim=1))
+
+                c_confidences, c_predictions = torch.max(softmax_prob, dim=1)
+                c_predictions=c_predictions.tolist()
+                c_confidences=c_confidences.tolist()
+                c_ece = expected_calibration_error(c_confidences, c_predictions, labels_list, num_bins=15)
+                print('ece', ece, 't_ece', t_ece, 'c_ece', c_ece)
+                c_accuracy=accuracy_score(labels_list, c_predictions)
+                print('accuracy',accuracy,'t_accuracy',t_accuracy,'c_accuracy',c_accuracy)
+
+                ood_logits,ood_mars=SPC_evaluate(net, SPC_model, ood_test_loader, model_to_num_dim[args.model], num_classes, device)
+                (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc_logits(mars, ood_mars, self_consistency, device)
+                (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc_logits(logits, ood_logits, entropy, device)
+
+                labels_array = np.array(labels_list)
+                pred_array = np.array(predictions)
+                correct_mask = labels_array == pred_array
+                mars_right = [m[correct_mask] for m in mars]
+                mars_wrong = [m[~correct_mask] for m in mars]
+                (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(mars_right, mars_wrong, self_consistency, device)
+                (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc = get_roc_auc_logits(mars_right[0], mars_wrong[0], entropy, device)
+
+                if args.dataset == 'imagenet':
+                    adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                                   imagesize=model_to_input_dim[args.model],
+                                                                                   pin_memory=args.gpu)
+                else:
+                    adv_test_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,batch_size=args.batch_size)
+                (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,adv_confidences,) = test_classification_net(net, adv_test_loader, device)
+                adv_logits, adv_mars = SPC_evaluate(net, SPC_model, adv_test_loader,model_to_num_dim[args.model], num_classes, device)
+                (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(mars, adv_mars,self_consistency, device)
+                (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits, adv_logits, entropy, device)
+
+                if args.sample_noise:
+                    adv_eps = np.linspace(0, 0.4, 9)
+                    print(adv_eps)
+                    for idx_ep, ep in enumerate(adv_eps):
+                        adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=ep,
+                                                                   batch_size=args.batch_size)
+                        (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                         adv_confidences,) = test_classification_net(net, adv_loader, device)
+                        adv_logits, adv_mars=SPC_evaluate(net, SPC_model, adv_loader, model_to_num_dim[args.model], num_classes, device)
+                        uncertainties = self_consistency(adv_mars).detach().cpu().numpy()
+                        quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                        quantiles = np.delete(quantiles, 0)
+                        unc_list = []
+                        accuracy_list = []
+                        for threshold in quantiles:
+                            cer_indices = (uncertainties < threshold)
+                            unc_indices = ~cer_indices
+                            labels_list = np.array(adv_labels_list)
+                            targets_cer = labels_list[cer_indices]
+                            predictions = np.array(adv_predictions)
+                            pred_cer = predictions[cer_indices]
+                            targets_unc = labels_list[unc_indices]
+                            pred_unc = predictions[unc_indices]
+                            cer_right = np.sum(targets_cer == pred_cer)
+                            cer = len(targets_cer)
+                            unc_right = np.sum(targets_unc == pred_unc)
+                            unc = len(targets_unc)
+                            accuracy_cer = cer_right / cer
+                            accuracy_unc = unc_right / unc
+                            unc_list.append(threshold)
+                            accuracy_list.append(accuracy_cer)
+                            print('ACC:', accuracy_cer, accuracy_unc)
+                        from scipy.stats import spearmanr
+                        Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                        print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                        adv_unc[i][idx_ep] = uncertainties.mean()
+                        adv_acc[i][idx_ep] = adv_accuracy
+
+                        (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc_logits(mars, adv_mars, self_consistency,device)
+                        (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc_logits(logits, adv_logits, entropy,device)
+                        print('adv_m1_auroc,adv_m2_auroc',adv_m1_auroc,adv_m2_auroc)
+
+
+                t_m1_auroc = ood_m1_auroc
+                t_m1_auprc = ood_m1_auprc
+                t_m2_auroc = ood_m2_auroc
+                t_m2_auprc = ood_m2_auprc
+
+            else:
+                # Evaluate a normal Softmax model
+                print("Softmax Model")
+                (_, _, _), (_, _, _), ood_m1_auroc, ood_m1_auprc = get_roc_auc(net, test_loader, ood_test_loader, entropy, device)
+                (_, _, _), (_, _, _), ood_m2_auroc, ood_m2_auprc = get_roc_auc(net, test_loader, ood_test_loader, logsumexp, device, confidence=True)
+
+                (_, _, _), (_, _, _), t_m1_auroc, t_m1_auprc = get_roc_auc(temp_scaled_net, test_loader, ood_test_loader, entropy, device)
+                (_, _, _), (_, _, _), t_m2_auroc, t_m2_auprc = get_roc_auc(temp_scaled_net, test_loader, ood_test_loader, logsumexp, device, confidence=True)
+
+                labels_array = np.array(labels_list)
+                pred_array = np.array(predictions)
+                correct_mask = labels_array == pred_array
+                logits, _ = get_logits_labels(net, test_loader, device)
+                logits_right = logits[correct_mask]
+                logits_wrong = logits[~correct_mask]
+                (_, _, _), (_, _, _), err_m1_auroc, err_m1_auprc = get_roc_auc_logits(logits_right, logits_wrong, entropy, device)
+                (_, _, _), (_, _, _), err_m2_auroc, err_m2_auprc = get_roc_auc_logits(logits_right, logits_wrong, logsumexp, device, confidence=True)
+
+                if args.dataset == 'imagenet':
+                    adv_test_loader = dataset_loader['imagenet_a'].get_test_loader(batch_size=args.batch_size,
+                                                                                   imagesize=model_to_input_dim[args.model],
+                                                                                   pin_memory=args.gpu)
+                else:
+                    adv_test_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep,batch_size=args.batch_size)
+                (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions, adv_confidences,) = test_classification_net(net, adv_test_loader, device)
+                adv_logits, _ = get_logits_labels(net, adv_test_loader, device)
+                (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc(net, test_loader, adv_test_loader, entropy, device)
+                (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc(net, test_loader, adv_test_loader, logsumexp, device, confidence=True)
+
+                if args.sample_noise:
+                    adv_eps = np.linspace(0, 0.4, 9)
+                    for idx_ep, ep in enumerate(adv_eps):
+                        adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=ep,
+                                                                   batch_size=args.batch_size)
+                        (adv_conf_matrix, adv_accuracy, adv_labels_list, adv_predictions,
+                         adv_confidences,) = test_classification_net(net, adv_loader, device)
+                        adv_logits, _ = get_logits_labels(net, adv_loader, device)
+                        uncertainties = entropy(adv_logits).detach().cpu().numpy()
+                        quantiles = np.quantile(uncertainties, np.linspace(0, 1, 10))
+                        quantiles = np.delete(quantiles, 0)
+                        unc_list = []
+                        accuracy_list = []
+                        for threshold in quantiles:
+                            cer_indices = (uncertainties < threshold)
+                            unc_indices = ~cer_indices
+                            labels_list = np.array(adv_labels_list)
+                            targets_cer = labels_list[cer_indices]
+                            predictions = np.array(adv_predictions)
+                            pred_cer = predictions[cer_indices]
+                            targets_unc = labels_list[unc_indices]
+                            pred_unc = predictions[unc_indices]
+                            cer_right = np.sum(targets_cer == pred_cer)
+                            cer = len(targets_cer)
+                            unc_right = np.sum(targets_unc == pred_unc)
+                            unc = len(targets_unc)
+                            accuracy_cer = cer_right / cer
+                            accuracy_unc = unc_right / unc
+                            unc_list.append(threshold)
+                            accuracy_list.append(accuracy_cer)
+                            print('ACC:', accuracy_cer, accuracy_unc)
+                        from scipy.stats import spearmanr
+
+                        Spearman_acc, p_acc = spearmanr(unc_list, accuracy_list)
+                        print("Spearman correlation:", Spearman_acc, "mean uncertainties:", uncertainties.mean())
+                        adv_unc[i][idx_ep] = uncertainties.mean()
+                        adv_acc[i][idx_ep] = adv_accuracy
+
+                        (_, _, _), (_, _, _), adv_m1_auroc, adv_m1_auprc = get_roc_auc(net, test_loader, adv_loader, entropy, device)
+                        (_, _, _), (_, _, _), adv_m2_auroc, adv_m2_auprc = get_roc_auc(net, test_loader, adv_loader, logsumexp, device, confidence=True)
+                        print('adv_m1_auroc,adv_m2_auroc', adv_m1_auroc, adv_m2_auroc)
+
+
+        accuracies.append(accuracy)
+        if (args.model_type == "spc" or args.model_type == "joint"):
+            c_accuracies.append(c_accuracy)
+        else:
+            c_accuracies.append(t_accuracy)
+
+
+        # Pre-temperature results
+        eces.append(ece)
+        ood_m1_aurocs.append(ood_m1_auroc)
+        ood_m1_auprcs.append(ood_m1_auprc)
+        ood_m2_aurocs.append(ood_m2_auroc)
+        ood_m2_auprcs.append(ood_m2_auprc)
+
+        err_m1_aurocs.append(err_m1_auroc)
+        err_m1_auprcs.append(err_m1_auprc)
+        err_m2_aurocs.append(err_m2_auroc)
+        err_m2_auprcs.append(err_m2_auprc)
+
+        adv_m1_aurocs.append(adv_m1_auroc)
+        adv_m1_auprcs.append(adv_m1_auprc)
+        adv_m2_aurocs.append(adv_m2_auroc)
+        adv_m2_auprcs.append(adv_m2_auprc)
+
+        # Post-temperature results
+        t_eces.append(t_ece)
+        t_m1_aurocs.append(t_m1_auroc)
+        t_m1_auprcs.append(t_m1_auprc)
+        t_m2_aurocs.append(t_m2_auroc)
+        t_m2_auprcs.append(t_m2_auprc)
+
+        if (args.model_type == "spc" or args.model_type == "joint"):
+            c_eces.append(c_ece)
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+    if args.sample_noise:
+        adv_unc_norm = (adv_unc - adv_unc.min(axis=1, keepdims=True)) / \
+                       (adv_unc.max(axis=1, keepdims=True) - adv_unc.min(axis=1, keepdims=True) + 1e-8)
+        mean_unc = np.mean(adv_unc_norm, axis=0)
+        std_unc = np.std(adv_unc_norm, axis=0)
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(mean_unc, label='Uncertainty', color='orange')
+        plt.fill_between(range(len(mean_unc)),
+                         mean_unc - std_unc,
+                         mean_unc + std_unc,
+                         color='orange', alpha=0.3, label="±1 Std Dev")
+        # plt.legend()
+        # plt.title('Uncertainty Across Multiple Runs')
+        plt.xlabel('Noise')
+        plt.ylabel('Uncertainty')
+        plt.savefig("adv_unc_"
+            + model_save_name(args.model, args.sn, args.mod, args.coeff, args.seed)
+            + "_"
+            + args.model_type
+            + "_"
+            + args.dataset
+            + ".png")
+        plt.show()
+        plt.close()
+
+
+        mean_acc = np.mean(adv_acc, axis=0)
+        std_acc = np.std(adv_acc, axis=0)
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(mean_acc, label='Accuracy', color='red')
+        plt.fill_between(range(len(mean_acc)),
+                         mean_acc - std_acc,
+                         mean_acc + std_acc,
+                         color='red', alpha=0.3, label="±1 Std Dev")
+        # plt.legend()
+        # plt.title('Uncertainty Across Multiple Runs')
+        plt.xlabel('Noise')
+        plt.ylabel('Accuracy')
+        plt.savefig("adv_acc_"
+            + model_save_name(args.model, args.sn, args.mod, args.coeff, args.seed)
+            + "_"
+            + args.model_type
+            + "_"
+            + args.dataset
+            + ".png")
+        plt.show()
+        plt.close()
+
+        save_dir = "curve_data"
+        os.makedirs(save_dir, exist_ok=True)
+
+        if args.sn:
+            prefix = f"{args.dataset}_{args.model_type}_{args.model}_SN"
+        else:
+            prefix = f"{args.dataset}_{args.model_type}_{args.model}"
+
+
+    accuracy_tensor = torch.tensor(accuracies)
+    c_accuracy_tensor = torch.tensor(c_accuracies)
+    ece_tensor = torch.tensor(eces)
+    ood_m1_auroc_tensor = torch.tensor(ood_m1_aurocs)
+    m1_auprc_tensor = torch.tensor(ood_m1_auprcs)
+    ood_m2_auroc_tensor = torch.tensor(ood_m2_aurocs)
+    ood_m2_auprc_tensor = torch.tensor(ood_m2_auprcs)
+
+    err_m1_auroc_tensor = torch.tensor(err_m1_aurocs)
+    err_m1_auprc_tensor = torch.tensor(err_m1_auprcs)
+    err_m2_auroc_tensor = torch.tensor(err_m2_aurocs)
+    err_m2_auprc_tensor = torch.tensor(err_m2_auprcs)
+
+    adv_m1_auroc_tensor = torch.tensor(adv_m1_aurocs)
+    adv_m1_auprc_tensor = torch.tensor(adv_m1_auprcs)
+    adv_m2_auroc_tensor = torch.tensor(adv_m2_aurocs)
+    adv_m2_auprc_tensor = torch.tensor(adv_m2_auprcs)
+
+    t_ece_tensor = torch.tensor(t_eces)
+    t_m1_auroc_tensor = torch.tensor(t_m1_aurocs)
+    t_m1_auprc_tensor = torch.tensor(t_m1_auprcs)
+    t_m2_auroc_tensor = torch.tensor(t_m2_aurocs)
+    t_m2_auprc_tensor = torch.tensor(t_m2_auprcs)
+
+    c_ece_tensor = torch.tensor(c_eces)
+
+    mean_accuracy = torch.mean(accuracy_tensor)
+    mean_c_accuracy = torch.mean(c_accuracy_tensor)
+    mean_ece = torch.mean(ece_tensor)
+    mean_ood_m1_auroc = torch.mean(ood_m1_auroc_tensor)
+    mean_m1_auprc = torch.mean(m1_auprc_tensor)
+    mean_m2_auroc = torch.mean(ood_m2_auroc_tensor)
+    mean_m2_auprc = torch.mean(ood_m2_auprc_tensor)
+
+    mean_err_m1_auroc = torch.mean(err_m1_auroc_tensor)
+    mean_err_m1_auprc = torch.mean(err_m1_auprc_tensor)
+    mean_err_m2_auroc = torch.mean(err_m2_auroc_tensor)
+    mean_err_m2_auprc = torch.mean(err_m2_auprc_tensor)
+
+    mean_adv_m1_auroc = torch.mean(adv_m1_auroc_tensor)
+    mean_adv_m1_auprc = torch.mean(adv_m1_auprc_tensor)
+    mean_adv_m2_auroc = torch.mean(adv_m2_auroc_tensor)
+    mean_adv_m2_auprc = torch.mean(adv_m2_auprc_tensor)
+
+    mean_t_ece = torch.mean(t_ece_tensor)
+    mean_t_m1_auroc = torch.mean(t_m1_auroc_tensor)
+    mean_t_m1_auprc = torch.mean(t_m1_auprc_tensor)
+    mean_t_m2_auroc = torch.mean(t_m2_auroc_tensor)
+    mean_t_m2_auprc = torch.mean(t_m2_auprc_tensor)
+
+    mean_c_ece = torch.mean(c_ece_tensor)
+
+    std_accuracy = torch.std(accuracy_tensor) / math.sqrt(accuracy_tensor.shape[0])
+    std_c_accuracy = torch.std(c_accuracy_tensor) / math.sqrt(c_accuracy_tensor.shape[0])
+    std_ece = torch.std(ece_tensor) / math.sqrt(ece_tensor.shape[0])
+    std_ood_m1_auroc = torch.std(ood_m1_auroc_tensor) / math.sqrt(ood_m1_auroc_tensor.shape[0])
+    std_m1_auprc = torch.std(m1_auprc_tensor) / math.sqrt(m1_auprc_tensor.shape[0])
+    std_m2_auroc = torch.std(ood_m2_auroc_tensor) / math.sqrt(ood_m2_auroc_tensor.shape[0])
+    std_m2_auprc = torch.std(ood_m2_auprc_tensor) / math.sqrt(ood_m2_auprc_tensor.shape[0])
+    std_err_m1_auroc = torch.std(err_m1_auroc_tensor) / math.sqrt(err_m1_auroc_tensor.shape[0])
+    std_err_m1_auprc = torch.std(err_m1_auprc_tensor) / math.sqrt(err_m1_auprc_tensor.shape[0])
+    std_err_m2_auroc = torch.std(err_m2_auroc_tensor) / math.sqrt(err_m2_auroc_tensor.shape[0])
+    std_err_m2_auprc = torch.std(err_m2_auprc_tensor) / math.sqrt(err_m2_auprc_tensor.shape[0])
+    std_adv_m1_auroc = torch.std(adv_m1_auroc_tensor) / math.sqrt(adv_m1_auroc_tensor.shape[0])
+    std_adv_m1_auprc = torch.std(adv_m1_auprc_tensor) / math.sqrt(adv_m1_auprc_tensor.shape[0])
+    std_adv_m2_auroc = torch.std(adv_m2_auroc_tensor) / math.sqrt(adv_m2_auroc_tensor.shape[0])
+    std_adv_m2_auprc = torch.std(adv_m2_auprc_tensor) / math.sqrt(adv_m2_auprc_tensor.shape[0])
+
+    std_t_ece = torch.std(t_ece_tensor) / math.sqrt(t_ece_tensor.shape[0])
+    std_t_m1_auroc = torch.std(t_m1_auroc_tensor) / math.sqrt(t_m1_auroc_tensor.shape[0])
+    std_t_m1_auprc = torch.std(t_m1_auprc_tensor) / math.sqrt(t_m1_auprc_tensor.shape[0])
+    std_t_m2_auroc = torch.std(t_m2_auroc_tensor) / math.sqrt(t_m2_auroc_tensor.shape[0])
+    std_t_m2_auprc = torch.std(t_m2_auprc_tensor) / math.sqrt(t_m2_auprc_tensor.shape[0])
+
+    std_c_ece = torch.std(c_ece_tensor) / math.sqrt(c_ece_tensor.shape[0])
+
+    res_dict = {}
+    res_dict["mean"] = {}
+    res_dict["mean"]["accuracy"] = mean_accuracy.item()
+    res_dict["mean"]["ece"] = mean_ece.item()
+    res_dict["mean"]["ood_m1_auroc"] = mean_ood_m1_auroc.item()
+    res_dict["mean"]["ood_m1_auprc"] = mean_m1_auprc.item()
+    res_dict["mean"]["ood_m2_auroc"] = mean_m2_auroc.item()
+    res_dict["mean"]["ood_m2_auprc"] = mean_m2_auprc.item()
+    res_dict["mean"]["t_ece"] = mean_t_ece.item()
+    res_dict["mean"]["t_m1_auroc"] = mean_t_m1_auroc.item()
+    res_dict["mean"]["t_m1_auprc"] = mean_t_m1_auprc.item()
+    res_dict["mean"]["t_m2_auroc"] = mean_t_m2_auroc.item()
+    res_dict["mean"]["t_m2_auprc"] = mean_t_m2_auprc.item()
+    res_dict["mean"]["c_ece"] = mean_c_ece.item()
+
+    res_dict["std"] = {}
+    res_dict["std"]["accuracy"] = std_accuracy.item()
+    res_dict["std"]["ece"] = std_ece.item()
+    res_dict["std"]["ood_m1_auroc"] = std_ood_m1_auroc.item()
+    res_dict["std"]["ood_m1_auprc"] = std_m1_auprc.item()
+    res_dict["std"]["ood_m2_auroc"] = std_m2_auroc.item()
+    res_dict["std"]["ood_m2_auprc"] = std_m2_auprc.item()
+    res_dict["std"]["t_ece"] = std_t_ece.item()
+    res_dict["std"]["t_m1_auroc"] = std_t_m1_auroc.item()
+    res_dict["std"]["t_m1_auprc"] = std_t_m1_auprc.item()
+    res_dict["std"]["t_m2_auroc"] = std_t_m2_auroc.item()
+    res_dict["std"]["t_m2_auprc"] = std_t_m2_auprc.item()
+    res_dict["std"]["c_ece"] = std_c_ece.item()
+
+    res_dict["values"] = {}
+    res_dict["values"]["accuracy"] = accuracies
+    res_dict["values"]["ece"] = eces
+    res_dict["values"]["ood_m1_auroc"] = ood_m1_aurocs
+    res_dict["values"]["ood_m1_auprc"] = ood_m1_auprcs
+    res_dict["values"]["ood_m2_auroc"] = ood_m2_aurocs
+    res_dict["values"]["ood_m2_auprc"] = ood_m2_auprcs
+    res_dict["values"]["t_ece"] = t_eces
+    res_dict["values"]["t_m1_auroc"] = t_m1_aurocs
+    res_dict["values"]["t_m1_auprc"] = t_m1_auprcs
+    res_dict["values"]["t_m2_auroc"] = t_m2_aurocs
+    res_dict["values"]["t_m2_auprc"] = t_m2_auprcs
+    res_dict["values"]["c_ece"] = c_eces
+
+    res_dict["info"] = vars(args)
+
+    print(f"{mean_accuracy.item() * 100:.2f} ± {std_accuracy.item() * 100:.2f}")
+    print(f"{mean_c_accuracy.item() * 100:.2f} ± {std_c_accuracy.item() * 100:.2f}")
+    print(f"{mean_ece.item()*100:.2f} ± {std_ece.item()*100:.2f}")
+    print(f"{mean_t_ece.item()*100:.2f} ± {std_t_ece.item()*100:.2f}")
+    print(f"{mean_c_ece.item() * 100:.2f} ± {std_c_ece.item() * 100:.2f}")
+    print(f"{mean_adv_m1_auroc.item()*100:.2f} ± {std_adv_m1_auroc.item()*100:.2f}")
+    print(f"{mean_err_m1_auroc.item()*100:.2f} ± {std_err_m1_auroc.item()*100:.2f}")
+    print(f"{mean_ood_m1_auroc.item()*100:.2f} ± {std_ood_m1_auroc.item()*100:.2f}")
+
+
+    with open(
+        "res_"
+        + model_save_name(args.model, args.sn, args.mod, args.coeff, args.seed)
+        + "_"
+        + args.model_type
+        + "_"
+        + args.dataset
+        + "_"
+        + args.ood_dataset
+        + ".json",
+        "w",
+    ) as f:
+        json.dump(res_dict, f)

SPC-UQ/Image_Classification/evaluate_laplace.py

@@ -0,0 +1,355 @@
+"""
+Script to evaluate the Laplace Approximation.
+"""
+
+import os
+import json
+import math
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import argparse
+import torch.backends.cudnn as cudnn
+import numpy as np
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+# Import data loaders and networks
+import data.ood_detection.cifar10 as cifar10
+import data.ood_detection.cifar100 as cifar100
+import data.ood_detection.svhn as svhn
+import data.ood_detection.imagenet as imagenet
+import data.ood_detection.tinyimagenet as tinyimagenet
+import data.ood_detection.imagenet_o as imagenet_o
+import data.ood_detection.imagenet_a as imagenet_a
+import data.ood_detection.ood_union as ood_union
+
+from net.resnet import resnet50
+from net.resnet_edl import resnet50_edl
+from net.wide_resnet import wrn
+from net.wide_resnet_edl import wrn_edl
+from net.vgg import vgg16
+from net.vgg_edl import vgg16_edl
+from net.imagenet_wide import imagenet_wide
+from net.imagenet_vgg import imagenet_vgg16
+from net.imagenet_vit import imagenet_vit
+
+from metrics.classification_metrics import (
+    test_classification_net,
+    test_classification_net_logits,
+    test_classification_net_ensemble,
+    test_classification_net_edl,
+    create_adversarial_dataloader
+)
+from metrics.calibration_metrics import expected_calibration_error
+
+from utils.gmm_utils import get_embeddings, gmm_evaluate, gmm_fit
+from utils.ensemble_utils import load_ensemble, ensemble_forward_pass
+from utils.eval_utils import model_load_name
+from utils.train_utils import model_save_name
+from utils.args import laplace_eval_args
+
+from laplace import Laplace
+from sklearn import metrics as M
+from laplace.curvature import AsdlGGN, AsdlEF, BackPackGGN, BackPackEF
+
+import warnings
+warnings.filterwarnings('ignore')
+
+# Dataset mapping and config
+DATASET_NUM_CLASSES = {
+    "cifar10": 10, "cifar100": 100, "svhn": 10, "imagenet": 1000,
+    "tinyimagenet": 200, "imagenet_o": 200, "imagenet_a": 200
+}
+
+DATASET_LOADER = {
+    "cifar10": cifar10, "cifar100": cifar100, "svhn": svhn, "imagenet": imagenet,
+    "tinyimagenet": tinyimagenet, "imagenet_o": imagenet_o, "imagenet_a": imagenet_a
+}
+
+MODELS = {
+    "resnet50": resnet50, "resnet50_edl": resnet50_edl,
+    "wide_resnet": wrn, "wide_resnet_edl": wrn_edl,
+    "vgg16": vgg16, "vgg16_edl": vgg16_edl,
+    "imagenet_wide": imagenet_wide
+}
+
+MODEL_TO_NUM_DIM = {
+    "resnet50": 2048, "resnet50_edl": 2048, "wide_resnet": 640, "wide_resnet_edl": 640,
+    "vgg16": 512, "vgg16_edl": 512, "imagenet_wide": 2048,
+    "imagenet_vgg16": 4096, "imagenet_vit": 768
+}
+
+MODEL_TO_INPUT_DIM = {
+    "resnet50": 32, "resnet50_edl": 32, "wide_resnet": 32, "wide_resnet_edl": 32,
+    "vgg16": 32, "vgg16_edl": 32, "imagenet_wide": 224,
+    "imagenet_vgg16": 224, "imagenet_vit": 224
+}
+
+MODEL_TO_LAST_LAYER = {
+    "resnet50": "module.fc", "wide_resnet": "module.linear", "vgg16": "module.classifier",
+    "imagenet_wide": "module.linear",
+    "imagenet_vgg16": "module.classifier", "imagenet_vit": "module.linear"
+}
+
+def get_backend(backend, approx_type):
+    if backend == 'kazuki':
+        return AsdlGGN if approx_type == 'ggn' else AsdlEF
+    elif backend == 'backpack':
+        return BackPackGGN if approx_type == 'ggn' else BackPackEF
+    else:
+        raise ValueError(f"Unknown backend: {backend}")
+
+def get_cpu_memory_mb():
+    import psutil
+    process = psutil.Process(os.getpid())
+    mem_info = process.memory_info()
+    return mem_info.rss / 1024 ** 2
+
+def print_metrics(mean, std, name):
+    print(f"{name}: {mean * 100:.2f} ± {std * 100:.2f}")
+
+if __name__ == "__main__":
+
+    args = laplace_eval_args().parse_args()
+
+    # Set random seed and device
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Parsed args:", args)
+    print("Seed:", args.seed)
+
+    num_classes = DATASET_NUM_CLASSES[args.dataset]
+    test_loader = DATASET_LOADER[args.dataset].get_test_loader(
+        batch_size=args.batch_size, imagesize=MODEL_TO_INPUT_DIM[args.model], pin_memory=args.gpu
+    )
+    if args.ood_dataset == 'ood_union':
+        ood_test_loader = ood_union.get_combined_ood_test_loader(
+            batch_size=args.batch_size, sample_seed=args.seed,
+            imagesize=MODEL_TO_INPUT_DIM[args.model], pin_memory=args.gpu
+        )
+    else:
+        ood_test_loader = DATASET_LOADER[args.ood_dataset].get_test_loader(
+            batch_size=args.batch_size, imagesize=MODEL_TO_INPUT_DIM[args.model], pin_memory=args.gpu
+        )
+
+    # Prepare metric accumulators
+    accuracies, eces, ood_aurocs, err_aurocs, adv_aurocs = [], [], [], [], []
+    err_aurocs, adv_aurocs = [], []
+    adv_unc = np.zeros((args.runs, 9))
+    adv_acc = np.zeros((args.runs, 9))
+    adv_ep = 0.02
+
+    for i in range(args.runs):
+        # Load training/validation splits
+        train_loader, val_loader = DATASET_LOADER[args.dataset].get_train_valid_loader(
+            batch_size=args.batch_size, imagesize=MODEL_TO_INPUT_DIM[args.model], augment=args.data_aug,
+            val_seed=(args.seed + i), val_size=args.val_size, pin_memory=args.gpu
+        )
+        mixture_components = []
+        for model_idx in range(args.nr_components):
+            if args.dataset == 'imagenet':
+                net = MODELS[args.model](pretrained=True, num_classes=1000).cuda()
+                net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
+                cudnn.benchmark = True
+            else:
+                if args.val_size == 0.1 or not args.crossval:
+                    saved_model_name = os.path.join(
+                        args.load_loc, f"Run{i+1}",
+                        model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i) + "_350.model",
+                    )
+                else:
+                    saved_model_name = os.path.join(
+                        args.load_loc, f"Run{i+1}",
+                        model_load_name(args.model, args.sn, args.mod, args.coeff, args.seed, i)
+                        + f"_350_0{int(args.val_size * 10)}.model"
+                    )
+                print('Loading:', saved_model_name)
+                net = MODELS[args.model](
+                    spectral_normalization=args.sn, mod=args.mod, coeff=args.coeff, num_classes=num_classes, temp=1.0
+                )
+                if args.gpu:
+                    net.cuda()
+                    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
+                    cudnn.benchmark = True
+                net.load_state_dict(torch.load(str(saved_model_name)))
+
+            # Laplace backend and fit
+            args.prior_precision = 1.0 if isinstance(args.prior_precision, float) else torch.load(args.prior_precision, map_location=device)
+            Backend = get_backend(args.backend, args.approx_type)
+            args.last_layer_name = MODEL_TO_LAST_LAYER[args.model]
+            optional_args = {"last_layer_name": args.last_layer_name} if args.subset_of_weights == 'last_layer' else {}
+
+            print('Fitting Laplace approximation...')
+            model = Laplace(
+                net, args.likelihood, subset_of_weights=args.subset_of_weights,
+                hessian_structure=args.hessian_structure, prior_precision=args.prior_precision,
+                temperature=args.temperature, backend=Backend, **optional_args
+            )
+            model.fit(val_loader if args.crossval else train_loader)
+
+            # Optional: Optimize prior precision
+            if (args.optimize_prior_precision is not None) and (args.method == 'laplace'):
+                n = model.n_params if args.prior_structure == 'all' else model.n_layers
+                prior_precision = args.prior_precision * torch.ones(n, device=device)
+                print('Optimizing prior precision...')
+                model.optimize_prior_precision(
+                    method=args.optimize_prior_precision, init_prior_prec=prior_precision,
+                    val_loader=val_loader, pred_type=args.pred_type, link_approx=args.link_approx,
+                    n_samples=args.n_samples, verbose=(args.prior_structure == 'scalar')
+                )
+            mixture_components.append(model)
+
+        model = mixture_components[0]
+        loss_fn = nn.NLLLoss()
+
+        # Evaluate ID data
+        id_y_true, id_y_prob = [], []
+        for data in tqdm(test_loader, desc='Evaluating ID data'):
+            x, y = data[0].to(device), data[1].to(device)
+            id_y_true.append(y.cpu())
+            y_prob = model(x, pred_type=args.pred_type, link_approx=args.link_approx, n_samples=args.n_samples)
+            id_y_prob.append(y_prob.cpu())
+        id_y_prob = torch.cat(id_y_prob, dim=0)
+        id_y_true = torch.cat(id_y_true, dim=0)
+        c, preds = torch.max(id_y_prob, 1)
+        metrics = {}
+        metrics['conf'] = c.mean().item()
+        metrics['nll'] = loss_fn(id_y_prob.log(), id_y_true).item()
+        metrics['acc'] = (id_y_true == preds).float().mean().item()
+        accuracy = metrics['acc']
+        id_confidences = id_y_prob.max(dim=1)[0].numpy()
+        ece = expected_calibration_error(id_confidences, preds.numpy(), id_y_true.numpy(), num_bins=15)
+        t_ece = ece
+        metrics['ece'] = ece
+        print(metrics)
+
+        # Evaluate OOD data
+        ood_y_true, ood_y_prob = [], []
+        for data in tqdm(ood_test_loader, desc='Evaluating OOD data'):
+            x, y = data[0].to(device), data[1].to(device)
+            ood_y_true.append(y.cpu())
+            y_prob = model(x, pred_type=args.pred_type, link_approx=args.link_approx, n_samples=args.n_samples)
+            ood_y_prob.append(y_prob.cpu())
+        ood_y_prob = torch.cat(ood_y_prob, dim=0)
+        ood_y_true = torch.cat(ood_y_true, dim=0)
+        ood_confidences = ood_y_prob.max(dim=1)[0].numpy()
+
+        # OOD AUROC/AUPRC metrics
+        bin_labels = np.concatenate([
+            np.zeros(id_confidences.shape[0]),
+            np.ones(ood_confidences.shape[0])
+        ])
+        scores = np.concatenate([id_confidences, ood_confidences])
+        fpr, tpr, thresholds = M.roc_curve(bin_labels, scores)
+        precision, recall, prc_thresholds = M.precision_recall_curve(bin_labels, scores)
+        ood_auroc = M.roc_auc_score(bin_labels, scores)
+        auprc = M.average_precision_score(bin_labels, scores)
+        print(f"OOD AUROC: {ood_auroc:.4f}, AUPRC: {auprc:.4f}")
+
+        # Error AUROC/AUPRC (in-distribution: correct vs incorrect)
+        labels_array = np.array(id_y_true)
+        pred_array = np.array(preds)
+        correct_mask = labels_array == pred_array
+        confidences_right = id_confidences[correct_mask]
+        confidences_wrong = id_confidences[~correct_mask]
+        bin_labels = np.concatenate([
+            np.zeros(confidences_right.shape[0]),
+            np.ones(confidences_wrong.shape[0])
+        ])
+        scores = np.concatenate([confidences_right, confidences_wrong])
+        err_auroc = M.roc_auc_score(bin_labels, scores)
+        err_auprc = M.average_precision_score(bin_labels, scores)
+        print(f"Error AUROC: {err_auroc:.4f}, AUPRC: {err_auprc:.4f}")
+
+        # Adversarial robustness
+        adv_loader = create_adversarial_dataloader(net, test_loader, device, epsilon=adv_ep, batch_size=args.batch_size)
+        adv_y_prob, adv_y_true = [], []
+        for data in tqdm(adv_loader, desc='Adversarial evaluation'):
+            x, y = data[0].to(device), data[1].to(device)
+            y_prob = model(x, pred_type=args.pred_type, link_approx=args.link_approx, n_samples=args.n_samples)
+            adv_y_true.append(y.cpu())
+            adv_y_prob.append(y_prob.cpu())
+        adv_y_prob = torch.cat(adv_y_prob, dim=0)
+        adv_y_true = torch.cat(adv_y_true, dim=0).numpy()
+        _, adv_predictions = torch.max(adv_y_prob, 1)
+        adv_accuracy = (adv_y_true == adv_predictions).mean()
+        adv_confidences = adv_y_prob.max(dim=1)[0].numpy()
+        bin_labels = np.concatenate([
+            np.zeros(id_confidences.shape[0]),
+            np.ones(adv_confidences.shape[0])
+        ])
+        adv_scores = np.concatenate([id_confidences, adv_confidences])
+        adv_auroc = M.roc_auc_score(bin_labels, adv_scores)
+        adv_auprc = M.average_precision_score(bin_labels, adv_scores)
+        print(f"Adversarial AUROC: {adv_auroc:.4f}, AUPRC: {adv_auprc:.4f}")
+
+        # If sample_noise: save/plot noise-uncertainty and accuracy curves
+        if args.sample_noise:
+            adv_eps = np.linspace(0, 0.4, 9)
+            for idx_ep, ep in enumerate(adv_eps):
+                adv_loader = create_adversarial_dataloader(
+                    net, test_loader, device, epsilon=ep, batch_size=args.batch_size
+                )
+                adv_y_prob, adv_y_true = [], []
+                for data in tqdm(adv_loader, desc=f"Adv evaluation ep={ep:.2f}"):
+                    x, y = data[0].to(device), data[1].to(device)
+                    y_prob = model(x, pred_type=args.pred_type, link_approx=args.link_approx, n_samples=args.n_samples)
+                    adv_y_true.append(y.cpu())
+                    adv_y_prob.append(y_prob.cpu())
+                adv_y_prob = torch.cat(adv_y_prob, dim=0)
+                adv_y_true = torch.cat(adv_y_true, dim=0).numpy()
+                _, predictions = torch.max(adv_y_prob, 1)
+                adv_accuracy = (adv_y_true == predictions).mean()
+                uncertainties = 1 - adv_y_prob.max(dim=1)[0].numpy()
+                adv_unc[i][idx_ep] = uncertainties.mean()
+                adv_acc[i][idx_ep] = adv_accuracy
+            # Save/plot uncertainty/accuracy curves as in your original
+
+        # Accumulate results
+        accuracies.append(accuracy)
+        eces.append(ece)
+        ood_aurocs.append(ood_auroc)
+        err_aurocs.append(err_auroc)
+        adv_aurocs.append(adv_auroc)
+
+        del model, mixture_components
+        torch.cuda.empty_cache()
+        import gc
+        gc.collect()
+
+    # Final result reporting and saving
+
+    def mean_std(x):
+        arr = torch.tensor(x)
+        return arr.mean().item(), arr.std().item() / math.sqrt(arr.shape[0])
+
+    # Print summary
+    print_metrics(*mean_std(accuracies), "Accuracy")
+    print_metrics(*mean_std(eces), "ECE")
+    print_metrics(*mean_std(adv_aurocs), "Adv AUROC")
+    print_metrics(*mean_std(err_aurocs), "Error AUROC")
+    print_metrics(*mean_std(ood_aurocs), "OOD AUROC")
+
+    # Store only required metrics
+    result_json = {}
+    for key, arr in [
+        ("accuracy", accuracies),
+        ("ece", eces),
+        ("adv_auroc", adv_aurocs),
+        ("err_auroc", err_aurocs),
+        ("ood_auroc", ood_aurocs)
+    ]:
+        mean, std = mean_std(arr)
+        result_json[key] = {
+            "mean": mean,
+            "std": std,
+            "values": [float(v) for v in arr]
+        }
+
+    result_file = (
+        "res_" + model_save_name(args.model, args.sn, args.mod, args.coeff, args.seed)
+        + "_laplace_" + args.dataset + "_" + args.ood_dataset + ".json"
+    )
+    with open(result_file, "w") as f:
+        json.dump(result_json, f, indent=2)

SPC-UQ/Image_Classification/metrics/calibration_metrics.py

@@ -0,0 +1,129 @@
+"""
+Metrics to measure calibration of a trained deep neural network.
+References:
+[1] C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger. On calibration of modern neural networks.
+    arXiv preprint arXiv:1706.04599, 2017.
+"""
+
+import math
+import torch
+import numpy as np
+from torch import nn
+from torch.nn import functional as F
+
+import matplotlib.pyplot as plt
+
+plt.rcParams.update({"font.size": 20})
+
+
+# Some keys used for the following dictionaries
+COUNT = "count"
+CONF = "conf"
+ACC = "acc"
+BIN_ACC = "bin_acc"
+BIN_CONF = "bin_conf"
+
+
+def _bin_initializer(num_bins=10):
+    bin_dict = {}
+    for i in range(num_bins):
+        bin_dict[i] = {}
+        bin_dict[i][COUNT] = 0
+        bin_dict[i][CONF] = 0
+        bin_dict[i][ACC] = 0
+        bin_dict[i][BIN_ACC] = 0
+        bin_dict[i][BIN_CONF] = 0
+
+    return bin_dict
+
+
+def _populate_bins(confs, preds, labels, num_bins=10):
+
+    bin_dict = _bin_initializer(num_bins)
+    num_test_samples = len(confs)
+
+    for i in range(0, num_test_samples):
+        confidence = confs[i]
+        prediction = preds[i]
+        label = labels[i]
+        # binn = int(math.ceil(((num_bins * confidence) - 1)))
+        binn = min(num_bins - 1, max(0, int(num_bins * confidence)))
+        # if binn>=num_bins:
+        #     binn=num_bins-1
+        bin_dict[binn][COUNT] = bin_dict[binn][COUNT] + 1
+        bin_dict[binn][CONF] = bin_dict[binn][CONF] + confidence
+        bin_dict[binn][ACC] = bin_dict[binn][ACC] + (1 if (label == prediction) else 0)
+
+    for binn in range(0, num_bins):
+        if bin_dict[binn][COUNT] == 0:
+            bin_dict[binn][BIN_ACC] = 0
+            bin_dict[binn][BIN_CONF] = 0
+        else:
+            bin_dict[binn][BIN_ACC] = float(bin_dict[binn][ACC]) / bin_dict[binn][COUNT]
+            bin_dict[binn][BIN_CONF] = bin_dict[binn][CONF] / float(bin_dict[binn][COUNT])
+    return bin_dict
+
+
+def expected_calibration_error(confs, preds, labels, num_bins=10):
+    bin_dict = _populate_bins(confs, preds, labels, num_bins)
+    num_samples = len(labels)
+    ece = 0
+    for i in range(num_bins):
+        bin_accuracy = bin_dict[i][BIN_ACC]
+        bin_confidence = bin_dict[i][BIN_CONF]
+        bin_count = bin_dict[i][COUNT]
+        ece += (float(bin_count) / num_samples) * abs(bin_accuracy - bin_confidence)
+    return ece
+
+
+# Calibration error scores in the form of loss metrics
+class ECELoss(nn.Module):
+    """
+    Compute ECE (Expected Calibration Error)
+    """
+
+    def __init__(self, n_bins=15):
+        super(ECELoss, self).__init__()
+        bin_boundaries = torch.linspace(0, 1, n_bins + 1)
+        self.bin_lowers = bin_boundaries[:-1]
+        self.bin_uppers = bin_boundaries[1:]
+
+    def forward(self, logits, labels):
+        softmaxes = F.softmax(logits, dim=1)
+        confidences, predictions = torch.max(softmaxes, 1)
+        accuracies = predictions.eq(labels)
+
+        ece = torch.zeros(1, device=logits.device)
+        for bin_lower, bin_upper in zip(self.bin_lowers, self.bin_uppers):
+            # Calculated |confidence - accuracy| in each bin
+            in_bin = confidences.gt(bin_lower.item()) * confidences.le(bin_upper.item())
+            prop_in_bin = in_bin.float().mean()
+            if prop_in_bin.item() > 0:
+                accuracy_in_bin = accuracies[in_bin].float().mean()
+                avg_confidence_in_bin = confidences[in_bin].mean()
+                ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
+
+        return ece
+
+
+# Methods for plotting reliability diagrams and bin-strength plots
+def reliability_plot(confs, preds, labels, num_bins=15, model_name='model'):
+    """
+    Method to draw a reliability plot from a model's predictions and confidences.
+    """
+    bin_dict = _populate_bins(confs, preds, labels, num_bins)
+    bns = [(i / float(num_bins)) for i in range(num_bins)]
+    y = []
+    for i in range(num_bins):
+        y.append(bin_dict[i][BIN_ACC])
+    plt.figure(figsize=(10, 8))  # width:20, height:3
+    plt.bar(bns, bns, align="edge", width=0.03, color="pink", label="Expected")
+    plt.bar(bns, y, align="edge", width=0.03, color="blue", alpha=0.5, label="Actual")
+    plt.ylabel("Accuracy", fontsize=30)
+    plt.xlabel("Confidence", fontsize=30)
+    plt.xticks(fontsize=30)
+    plt.yticks(fontsize=30)
+    plt.legend(fontsize=30, loc='upper left')
+    plt.savefig(f'./reliability_plot_{model_name}.pdf')
+    plt.savefig(f'./reliability_plot_{model_name}.png')
+    plt.show()

SPC-UQ/Image_Classification/metrics/classification_metrics.py

@@ -0,0 +1,211 @@
+"""
+Metrics to measure classification performance
+"""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader, TensorDataset
+
+from utils.ensemble_utils import ensemble_forward_pass
+
+from sklearn.metrics import accuracy_score
+from sklearn.metrics import confusion_matrix
+
+def evidential_loss(alpha, target, lambda_reg=0.001):
+    num_classes = alpha.shape[1]
+    target_one_hot = F.one_hot(target, num_classes=num_classes).float()
+    S = alpha.sum(dim=1, keepdim=True)
+
+    log_likelihood = torch.sum(target_one_hot * (torch.digamma(S) - torch.digamma(alpha)), dim=1)
+    kl_divergence = lambda_reg * torch.sum((alpha - 1) * (1 - target_one_hot), dim=1)
+
+    loss = log_likelihood + kl_divergence
+    return torch.mean(loss)
+
+def create_adversarial_dataloader(model, data_loader, device, epsilon=0.03, batch_size=32, edl=False, joint=False):
+    adv_examples = []
+    adv_labels = []
+
+    model.eval()
+    for data, label in data_loader:
+        data = data.to(device).detach().requires_grad_(True)
+        label = label.to(device)
+
+        model.zero_grad()
+        logit = model(data)
+        if edl:
+            loss = evidential_loss(logit, label)
+        if joint:
+            loss = F.cross_entropy(logit[0], label)
+        else:
+            loss = F.cross_entropy(logit, label)
+        loss.backward()
+
+        signed_grad = data.grad.sign()
+        data_adv = data + epsilon * signed_grad
+
+        adv_examples.append(data_adv.detach().cpu())
+        adv_labels.append(label.detach().cpu())
+
+    adv_examples = torch.cat(adv_examples, dim=0)
+    adv_labels = torch.cat(adv_labels, dim=0)
+
+    adv_dataset = TensorDataset(adv_examples, adv_labels)
+    adv_dataloader = DataLoader(adv_dataset, batch_size=batch_size, shuffle=False)
+
+    return adv_dataloader
+
+
+def get_logits_labels(model, data_loader, device):
+    """
+    Utility function to get logits and labels.
+    """
+    model.eval()
+    logits = []
+    labels = []
+    with torch.no_grad():
+        for data, label in data_loader:
+            data = data.to(device)
+            label = label.to(device)
+
+            logit = model(data)
+            logits.append(logit)
+            labels.append(label)
+    logits = torch.cat(logits, dim=0)
+    labels = torch.cat(labels, dim=0)
+    return logits, labels
+
+def get_logits_labels_uq(model, data_loader, device):
+    """
+    Utility function to get logits as a list: [pred_all, mar_all, mar_up_all, mar_down_all]
+    and labels as a single tensor.
+    """
+    model.eval()
+    preds = []
+    mars = []
+    mars_up = []
+    mars_down = []
+    labels = []
+
+    with torch.no_grad():
+        for data, label in data_loader:
+            data = data.to(device)
+            label = label.to(device)
+
+            pred, mar, mar_up, mar_down = model(data)
+            preds.append(pred)
+            mars.append(mar)
+            mars_up.append(mar_up)
+            mars_down.append(mar_down)
+            labels.append(label)
+
+    pred_all = torch.cat(preds, dim=0)
+    mar_all = torch.cat(mars, dim=0)
+    mar_up_all = torch.cat(mars_up, dim=0)
+    mar_down_all = torch.cat(mars_down, dim=0)
+    labels_all = torch.cat(labels, dim=0)
+
+    logits = [pred_all, mar_all, mar_up_all, mar_down_all]
+    return logits, labels_all
+
+
+def test_classification_net_softmax(softmax_prob, labels):
+    """
+    This function reports classification accuracy and confusion matrix given softmax vectors and
+    labels from a model.
+    """
+    labels_list = []
+    predictions_list = []
+    confidence_vals_list = []
+
+    confidence_vals, predictions = torch.max(softmax_prob, dim=1)
+    labels_list.extend(labels.cpu().numpy())
+    predictions_list.extend(predictions.cpu().numpy())
+    confidence_vals_list.extend(confidence_vals.detach().cpu().numpy())
+    accuracy = accuracy_score(labels_list, predictions_list)
+    return (
+        confusion_matrix(labels_list, predictions_list),
+        accuracy,
+        labels_list,
+        predictions_list,
+        confidence_vals_list,
+    )
+
+
+def test_classification_net_logits(logits, labels):
+    """
+    This function reports classification accuracy and confusion matrix given logits and labels
+    from a model.
+    """
+    softmax_prob = F.softmax(logits, dim=1)
+    return test_classification_net_softmax(softmax_prob, labels)
+
+
+def test_classification_net(model, data_loader, device):
+    """
+    This function reports classification accuracy and confusion matrix over a dataset.
+    """
+    logits, labels = get_logits_labels(model, data_loader, device)
+    return test_classification_net_logits(logits, labels)
+
+def test_classification_uq(model, data_loader, device):
+    """
+    This function reports classification accuracy and confusion matrix over a dataset.
+    """
+    logits, labels = get_logits_labels_uq(model, data_loader, device)
+    return test_classification_net_logits(logits[0], labels)
+
+def test_classification_net_ensemble(model_ensemble, data_loader, device):
+    """
+    This function reports classification accuracy and confusion matrix over a dataset
+    for a deep ensemble.
+    """
+    for model in model_ensemble:
+        model.eval()
+    softmax_prob = []
+    labels = []
+    with torch.no_grad():
+        for data, label in data_loader:
+            data = data.to(device)
+            label = label.to(device)
+
+            softmax, _, _ = ensemble_forward_pass(model_ensemble, data)
+            softmax_prob.append(softmax)
+            labels.append(label)
+    softmax_prob = torch.cat(softmax_prob, dim=0)
+    labels = torch.cat(labels, dim=0)
+
+    return test_classification_net_softmax(softmax_prob, labels)
+
+def test_classification_net_logits_edl(logits, labels):
+    """
+    This function reports classification accuracy and confusion matrix given softmax vectors and
+    labels from a model.
+    """
+    labels_list = []
+    predictions_list = []
+    confidence_vals_list = []
+
+    predicted_probs = logits / logits.sum(dim=1, keepdim=True)
+    confidence_vals, predictions = torch.max(predicted_probs, dim=1)
+    labels_list.extend(labels.cpu().numpy())
+    predictions_list.extend(predictions.cpu().numpy())
+    confidence_vals_list.extend(confidence_vals.cpu().numpy())
+    accuracy = accuracy_score(labels_list, predictions_list)
+    return (
+        confusion_matrix(labels_list, predictions_list),
+        accuracy,
+        labels_list,
+        predictions_list,
+        confidence_vals_list,
+    )
+
+def test_classification_net_edl(model, data_loader, device):
+    """
+    This function reports classification accuracy and confusion matrix over a dataset.
+    """
+    logits, labels = get_logits_labels(model, data_loader, device)
+    return test_classification_net_logits_edl(logits, labels)

SPC-UQ/Image_Classification/metrics/ood_metrics.py

@@ -0,0 +1,135 @@
+# Utility functions to get OOD detection ROC curves and AUROC scores
+# Ideally should be agnostic of model architectures
+
+import torch
+import torch.nn.functional as F
+from sklearn import metrics
+
+from utils.ensemble_utils import ensemble_forward_pass
+from metrics.classification_metrics import get_logits_labels
+from metrics.uncertainty_confidence import entropy, logsumexp, confidence, edl_unc
+
+
+def get_roc_auc(net, test_loader, ood_test_loader, uncertainty, device, confidence=False):
+    logits, _ = get_logits_labels(net, test_loader, device)
+    ood_logits, _ = get_logits_labels(net, ood_test_loader, device)
+
+    return get_roc_auc_logits(logits, ood_logits, uncertainty, device, confidence=confidence)
+
+
+def get_roc_auc_logits(logits, ood_logits, uncertainty, device, confidence=False):
+    uncertainties = uncertainty(logits)
+    ood_uncertainties = uncertainty(ood_logits)
+
+    # In-distribution
+    bin_labels = torch.zeros(uncertainties.shape[0]).to(device)
+    in_scores = uncertainties
+
+    # OOD
+    bin_labels = torch.cat((bin_labels, torch.ones(ood_uncertainties.shape[0]).to(device)))
+
+    if confidence:
+        bin_labels = 1 - bin_labels
+    ood_scores = ood_uncertainties  # entropy(ood_logits)
+    scores = torch.cat((in_scores, ood_scores))
+
+    fpr, tpr, thresholds = metrics.roc_curve(bin_labels.cpu().numpy(), scores.cpu().numpy())
+    precision, recall, prc_thresholds = metrics.precision_recall_curve(bin_labels.cpu().numpy(), scores.cpu().numpy())
+    auroc = metrics.roc_auc_score(bin_labels.cpu().numpy(), scores.cpu().numpy())
+    auprc = metrics.average_precision_score(bin_labels.cpu().numpy(), scores.cpu().numpy())
+
+    return (fpr, tpr, thresholds), (precision, recall, prc_thresholds), auroc, auprc
+
+
+def get_roc_auc_uncs(uncertainties, ood_uncertainties, device, confidence=False):
+    # In-distribution
+    bin_labels = torch.zeros(uncertainties.shape[0]).to(device)
+    in_scores = uncertainties
+
+    # OOD
+    bin_labels = torch.cat((bin_labels, torch.ones(ood_uncertainties.shape[0]).to(device)))
+
+    if confidence:
+        bin_labels = 1 - bin_labels
+    ood_scores = ood_uncertainties  # entropy(ood_logits)
+    scores = torch.cat((in_scores, ood_scores))
+
+    fpr, tpr, thresholds = metrics.roc_curve(bin_labels.detach().cpu().numpy(), scores.detach().cpu().numpy())
+    precision, recall, prc_thresholds = metrics.precision_recall_curve(bin_labels.detach().cpu().numpy(), scores.detach().cpu().numpy())
+    auroc = metrics.roc_auc_score(bin_labels.detach().cpu().numpy(), scores.detach().cpu().numpy())
+    auprc = metrics.average_precision_score(bin_labels.detach().cpu().numpy(), scores.detach().cpu().numpy())
+
+
+    return (fpr, tpr, thresholds), (precision, recall, prc_thresholds), auroc, auprc
+
+def get_roc_auc_ensemble(model_ensemble, test_loader, ood_test_loader, uncertainty, device):
+    bin_labels_uncertainties = None
+    uncertainties = None
+
+    for model in model_ensemble:
+        model.eval()
+
+    bin_labels_uncertainties = []
+    uncertainties = []
+    with torch.no_grad():
+        # Getting uncertainties for in-distribution data
+        for data, label in test_loader:
+            data = data.to(device)
+            label = label.to(device)
+
+            bin_label_uncertainty = torch.zeros(label.shape).to(device)
+            if uncertainty == "mutual_information":
+                net_output, _, unc = ensemble_forward_pass(model_ensemble, data)
+            else:
+                net_output, unc, _ = ensemble_forward_pass(model_ensemble, data)
+
+            bin_labels_uncertainties.append(bin_label_uncertainty)
+            uncertainties.append(unc)
+
+        # Getting entropies for OOD data
+        for data, label in ood_test_loader:
+            data = data.to(device)
+            label = label.to(device)
+
+            bin_label_uncertainty = torch.ones(label.shape).to(device)
+            if uncertainty == "mutual_information":
+                net_output, _, unc = ensemble_forward_pass(model_ensemble, data)
+            else:
+                net_output, unc, _ = ensemble_forward_pass(model_ensemble, data)
+
+            bin_labels_uncertainties.append(bin_label_uncertainty)
+            uncertainties.append(unc)
+
+        bin_labels_uncertainties = torch.cat(bin_labels_uncertainties)
+        uncertainties = torch.cat(uncertainties)
+
+    fpr, tpr, roc_thresholds = metrics.roc_curve(bin_labels_uncertainties.cpu().numpy(), uncertainties.cpu().numpy())
+    precision, recall, prc_thresholds = metrics.precision_recall_curve(
+        bin_labels_uncertainties.cpu().numpy(), uncertainties.cpu().numpy()
+    )
+    auroc = metrics.roc_auc_score(bin_labels_uncertainties.cpu().numpy(), uncertainties.cpu().numpy())
+    auprc = metrics.average_precision_score(bin_labels_uncertainties.cpu().numpy(), uncertainties.cpu().numpy())
+
+
+    return (fpr, tpr, roc_thresholds), (precision, recall, prc_thresholds), auroc, auprc
+
+
+def get_unc_ensemble(model_ensemble, test_loader, uncertainty, device):
+    for model in model_ensemble:
+        model.eval()
+
+    uncertainties = []
+    with torch.no_grad():
+        for data, label in test_loader:
+            data = data.to(device)
+
+            if uncertainty == "mutual_information":
+                net_output, _, unc = ensemble_forward_pass(model_ensemble, data)
+            else:
+                net_output, unc, _ = ensemble_forward_pass(model_ensemble, data)
+
+            uncertainties.append(unc)
+
+        uncertainties = torch.cat(uncertainties)
+
+    return uncertainties

SPC-UQ/Image_Classification/metrics/uncertainty_confidence.py

@@ -0,0 +1,67 @@
+"""
+Metrics measuring either uncertainty or confidence of a model.
+"""
+import torch
+import torch.nn.functional as F
+
+
+def entropy(logits):
+    p = F.softmax(logits, dim=1)
+    logp = F.log_softmax(logits, dim=1)
+    plogp = p * logp
+    entropy = -torch.sum(plogp, dim=1)
+    return entropy
+
+
+def logsumexp(logits):
+    return torch.logsumexp(logits, dim=1, keepdim=False)
+
+
+def confidence(logits):
+    p = F.softmax(logits, dim=1)
+    confidence, _ = torch.max(p, dim=1)
+    return confidence
+
+
+def entropy_prob(probs):
+    p = probs
+    eps = 1e-12
+    logp = torch.log(p + eps)
+    plogp = p * logp
+    entropy = -torch.sum(plogp, dim=1)
+    return entropy
+
+
+def mutual_information_prob(probs):
+    mean_output = torch.mean(probs, dim=0)
+    predictive_entropy = entropy_prob(mean_output)
+
+    # Computing expectation of entropies
+    p = probs
+    eps = 1e-12
+    logp = torch.log(p + eps)
+    plogp = p * logp
+    exp_entropies = torch.mean(-torch.sum(plogp, dim=2), dim=0)
+
+    # Computing mutual information
+    mi = predictive_entropy - exp_entropies
+    return mi
+
+def self_consistency(mars):
+    logits=mars[0]
+    logits = torch.nn.functional.softmax(logits, dim=1)
+    mar=mars[1].squeeze()
+    mar_up=1-logits
+    mar_down=logits
+    uncertainty = abs(2 * mar_up * mar_down - mar)   #(mar_up + mar_down)=1
+    uncertainty = torch.sum(uncertainty, dim=1)
+
+    return uncertainty
+
+def edl_unc(logits):
+    num_classes = logits.shape[1]
+    uncertainty = num_classes / logits.sum(dim=1)
+    return uncertainty
+
+def certificate(OCs):
+    return OCs

SPC-UQ/Image_Classification/net/imagenet_vgg.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+from torchvision.models import vgg16, VGG16_Weights
+
+
+class ImagenetVGG16(nn.Module):
+    """
+    VGG16 wrapper for ImageNet-like classification with:
+      - Optional pretrained backbone
+      - Optional feature freezing
+      - Temperature-scaled logits
+      - Exposes penultimate features via `self.feature` (detached)
+    """
+
+    def __init__(
+        self,
+        num_classes: int = 1000,
+        pretrained: bool = True,
+        temp: float = 1.0,
+        freeze_features: bool = False,
+    ) -> None:
+        super().__init__()
+
+        # Load base model (weights imply specific preprocessing; handle in dataloader)
+        base_model = vgg16(weights=VGG16_Weights.DEFAULT if pretrained else None)
+
+        # Convolutional feature extractor and avgpool
+        self.features: nn.Module = base_model.features
+        self.avgpool: nn.Module = base_model.avgpool
+
+        # Penultimate FC stack from original VGG16 (remove final classifier layer)
+        # VGG16 classifier:
+        # [Linear(25088->4096), ReLU, Dropout, Linear(4096->4096), ReLU, Dropout, Linear(4096->1000)]
+        self.fc_pre: nn.Sequential = nn.Sequential(*list(base_model.classifier.children())[:-1])
+
+        # New classification head
+        self.classifier: nn.Linear = nn.Linear(4096, num_classes)
+
+        # If using pretrained and num_classes matches 1000, copy the final layer weights
+        if pretrained and num_classes == 1000:
+            with torch.no_grad():
+                self.classifier.weight.copy_(base_model.classifier[-1].weight)
+                self.classifier.bias.copy_(base_model.classifier[-1].bias)
+
+        # Optional: freeze convolutional features (+ fc_pre) for linear probing
+        if freeze_features:
+            for p in self.features.parameters():
+                p.requires_grad = False
+            for p in self.fc_pre.parameters():
+                p.requires_grad = False
+
+        # Temperature (applied to logits at inference/training)
+        self.register_buffer("temperature", torch.tensor(float(temp)))
+        self.feature: Optional[torch.Tensor] = None  # cached detached penultimate features
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input batch of shape (B, 3, H, W)
+
+        Returns:
+            logits: Tensor of shape (B, num_classes)
+        """
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc_pre(x)
+
+        # Cache penultimate features (detached) for downstream use
+        self.feature = x.detach()
+
+        logits = self.classifier(x)
+        if self.temperature is not None and float(self.temperature) != 1.0:
+            logits = logits / self.temperature
+        return logits
+
+
+def imagenet_vgg16(
+    temp: float = 1.0,
+    pretrained: bool = True,
+    num_classes: int = 1000,
+    freeze_features: bool = False,
+    **kwargs,
+) -> ImagenetVGG16:
+    """
+    Factory function for ImagenetVGG16.
+
+    Args:
+        temp: Temperature applied to logits (T=1.0 disables scaling).
+        pretrained: Load pretrained ImageNet weights for the backbone.
+        num_classes: Output classes for the final classifier.
+        freeze_features: If True, freeze backbone + fc_pre for linear probing.
+        **kwargs: Forwarded to the ImagenetVGG16 init (future-proof).
+
+    Returns:
+        Initialized ImagenetVGG16 model.
+    """
+    return ImagenetVGG16(
+        num_classes=num_classes,
+        pretrained=pretrained,
+        temp=temp,
+        freeze_features=freeze_features,
+        **kwargs,
+    )

SPC-UQ/Image_Classification/net/imagenet_vit.py

@@ -0,0 +1,101 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+from torchvision.models import vit_b_16, ViT_B_16_Weights
+
+
+class ImagenetViT(nn.Module):
+    """
+    ViT-B/16 wrapper with:
+      - Optional pretrained backbone (torchvision)
+      - Proper CLS token + positional embeddings usage
+      - Temperature-scaled logits
+      - Exposed CLS feature via `self.feature` (detached)
+      - Optional backbone freezing for linear probing
+    """
+
+    def __init__(
+        self,
+        num_classes: int = 1000,
+        pretrained: bool = True,
+        temp: float = 1.0,
+        freeze_backbone: bool = False,
+    ) -> None:
+        super().__init__()
+        self.backbone = vit_b_16(weights=ViT_B_16_Weights.DEFAULT if pretrained else None)
+
+        # Hidden size & final norm
+        self.hidden_dim: int = self.backbone.hidden_dim
+        self.norm: nn.Module = self.backbone.encoder.ln  # final LayerNorm
+
+        # New classification head
+        self.head: nn.Linear = nn.Linear(self.hidden_dim, num_classes)
+
+        # If using pretrained and keeping 1000-class head, copy weights
+        if pretrained and num_classes == 1000:
+            with torch.no_grad():
+                src = self.backbone.heads.head
+                self.head.weight.copy_(src.weight)
+                self.head.bias.copy_(src.bias)
+
+        # Optionally freeze everything except the head
+        if freeze_backbone:
+            for p in self.backbone.parameters():
+                p.requires_grad = False
+
+        # Temperature buffer and CLS feature cache
+        self.register_buffer("temperature", torch.tensor(float(temp)))
+        self.feature: Optional[torch.Tensor] = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (B, 3, H, W)
+
+        Returns:
+            logits: (B, num_classes)
+        """
+        # Patchify + linear proj (handles resizing logic)
+        x = self.backbone._process_input(x)  # (B, N, hidden_dim)
+
+        # CLS token + positional embeddings (use pretrained params)
+        n = x.shape[0]
+        cls_token = self.backbone.class_token.expand(n, -1, -1)  # (B, 1, hidden_dim)
+        x = torch.cat((cls_token, x), dim=1)                      # (B, N+1, hidden_dim)
+
+        # Positional embeddings + dropout
+        x = x + self.backbone.encoder.pos_embedding               # (B, N+1, hidden_dim)
+        x = self.backbone.encoder.dropout(x)
+
+        # Transformer encoder (layers) + final norm
+        x = self.backbone.encoder.layers(x)
+        x = self.norm(x)
+
+        # CLS feature
+        cls = x[:, 0]                     # (B, hidden_dim)
+        self.feature = cls.detach()       # cache detached feature
+
+        # Head + optional temperature scaling
+        logits = self.head(cls)
+        if float(self.temperature) != 1.0:
+            logits = logits / self.temperature
+        return logits
+
+
+def imagenet_vit(
+    temp: float = 1.0,
+    pretrained: bool = True,
+    num_classes: int = 1000,
+    freeze_backbone: bool = False,
+    **kwargs,
+) -> ImagenetViT:
+    """
+    Factory for ImagenetViT.
+    """
+    return ImagenetViT(
+        num_classes=num_classes,
+        pretrained=pretrained,
+        temp=temp,
+        freeze_backbone=freeze_backbone,
+        **kwargs,
+    )

SPC-UQ/Image_Classification/net/imagenet_wide.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import wide_resnet50_2, Wide_ResNet50_2_Weights
+
+
+class ImagenetWideResNet(nn.Module):
+    def __init__(self, num_classes=1000, pretrained=True, temp=1.0):
+        super().__init__()
+
+        base_model = wide_resnet50_2(weights=Wide_ResNet50_2_Weights.DEFAULT if pretrained else None)
+
+        # Adapt to match your original WRN structure
+        self.features = nn.Sequential(
+            base_model.conv1,
+            base_model.bn1,
+            base_model.relu,
+            base_model.maxpool,
+            base_model.layer1,
+            base_model.layer2,
+            base_model.layer3,
+            base_model.layer4,
+            base_model.avgpool
+        )
+
+        self.linear = nn.Linear(2048, num_classes)
+        self.temp = temp
+        self.feature = None
+
+        if pretrained:
+            self.linear.load_state_dict(base_model.fc.state_dict())
+
+    def forward(self, x):
+        out = self.features(x)
+        out = torch.flatten(out, 1)
+        self.feature = out.clone().detach()
+        if self.temp == 1:
+            out = self.linear(out)
+        else:
+            out = self.linear(out) / self.temp
+        return out
+
+
+def imagenet_wide(temp=1.0, pretrained=True, **kwargs):
+    model = ImagenetWideResNet(pretrained=pretrained, temp=temp, **kwargs)
+    return model

SPC-UQ/Image_Classification/net/lenet.py

@@ -0,0 +1,37 @@
+"""Implementation of Lenet in pytorch.
+Refernece:
+[1] LeCun,  Y.,  Bottou,  L.,  Bengio,  Y.,  & Haffner,  P. (1998).
+    Gradient-based  learning  applied  to  document  recognition.
+    Proceedings of the IEEE, 86, 2278-2324.
+"""
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LeNet(nn.Module):
+    def __init__(self, num_classes, temp=1.0, mnist=True, **kwargs):
+        super(LeNet, self).__init__()
+        self.num_classes = num_classes
+        self.conv1 = nn.Conv2d(1 if mnist else 3, 6, 5)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(256, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, num_classes)
+        self.temp = temp
+        self.feature = None
+
+    def forward(self, x):
+        out = F.relu(self.conv1(x))
+        out = F.max_pool2d(out, 2)
+        out = F.relu(self.conv2(out))
+        out = F.max_pool2d(out, 2)
+        out = out.view(out.size(0), -1)
+        out = F.relu(self.fc1(out))
+        out = F.relu(self.fc2(out))
+        self.feature = out
+        out = self.fc3(out) / self.temp
+        return out
+
+
+def lenet(num_classes=10, temp=1.0, mnist=True, **kwargs):
+    return LeNet(num_classes=num_classes, temp=temp, mnist=True, **kwargs)

SPC-UQ/Image_Classification/net/resnet.py

@@ -0,0 +1,245 @@
+"""
+Pytorch implementation of ResNet models.
+Reference:
+[1] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR, 2016.
+"""
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+from net.spectral_normalization.spectral_norm_conv_inplace import spectral_norm_conv
+from net.spectral_normalization.spectral_norm_fc import spectral_norm_fc
+
+
+class AvgPoolShortCut(nn.Module):
+    def __init__(self, stride, out_c, in_c):
+        super(AvgPoolShortCut, self).__init__()
+        self.stride = stride
+        self.out_c = out_c
+        self.in_c = in_c
+
+    def forward(self, x):
+        if x.shape[2] % 2 != 0:
+            x = F.avg_pool2d(x, 1, self.stride)
+        else:
+            x = F.avg_pool2d(x, self.stride, self.stride)
+        pad = torch.zeros(x.shape[0], self.out_c - self.in_c, x.shape[2], x.shape[3], device=x.device,)
+        x = torch.cat((x, pad), dim=1)
+        return x
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, input_size, wrapped_conv, in_planes, planes, stride=1, mod=True):
+        super(BasicBlock, self).__init__()
+        self.conv1 = wrapped_conv(input_size, in_planes, planes, kernel_size=3, stride=stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = wrapped_conv(math.ceil(input_size / stride), planes, planes, kernel_size=3, stride=1)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.mod = mod
+        self.activation = F.leaky_relu if self.mod else F.relu
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            if mod:
+                self.shortcut = nn.Sequential(AvgPoolShortCut(stride, self.expansion * planes, in_planes))
+            else:
+                self.shortcut = nn.Sequential(
+                    wrapped_conv(input_size, in_planes, self.expansion * planes, kernel_size=1, stride=stride,),
+                    nn.BatchNorm2d(planes),
+                )
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = self.activation(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, input_size, wrapped_conv, in_planes, planes, stride=1, mod=True):
+        super(Bottleneck, self).__init__()
+        self.conv1 = wrapped_conv(input_size, in_planes, planes, kernel_size=1, stride=1)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = wrapped_conv(input_size, planes, planes, kernel_size=3, stride=stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = wrapped_conv(math.ceil(input_size / stride), planes, self.expansion * planes, kernel_size=1, stride=1)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+        self.mod = mod
+        self.activation = F.leaky_relu if self.mod else F.relu
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            if mod:
+                self.shortcut = nn.Sequential(AvgPoolShortCut(stride, self.expansion * planes, in_planes))
+            else:
+                self.shortcut = nn.Sequential(
+                    wrapped_conv(input_size, in_planes, self.expansion * planes, kernel_size=1, stride=stride,),
+                    nn.BatchNorm2d(self.expansion * planes),
+                )
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.activation(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = self.activation(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        block,
+        num_blocks,
+        num_classes=10,
+        temp=1.0,
+        spectral_normalization=True,
+        mod=True,
+        coeff=3,
+        n_power_iterations=1,
+        mnist=False,
+    ):
+        """
+        If the "mod" parameter is set to True, the architecture uses 2 modifications:
+        1. LeakyReLU instead of normal ReLU
+        2. Average Pooling on the residual connections.
+        """
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.mod = mod
+
+        def wrapped_conv(input_size, in_c, out_c, kernel_size, stride):
+            padding = 1 if kernel_size == 3 else 0
+
+            conv = nn.Conv2d(in_c, out_c, kernel_size, stride, padding, bias=False)
+
+            if not spectral_normalization:
+                return conv
+
+            # NOTE: Google uses the spectral_norm_fc in all cases
+            if kernel_size == 1:
+                # use spectral norm fc, because bound are tight for 1x1 convolutions
+                wrapped_conv = spectral_norm_fc(conv, coeff, n_power_iterations)
+            else:
+                # Otherwise use spectral norm conv, with loose bound
+                shapes = (in_c, input_size, input_size)
+                wrapped_conv = spectral_norm_conv(conv, coeff, shapes, n_power_iterations)
+
+            return wrapped_conv
+
+        self.wrapped_conv = wrapped_conv
+
+        self.bn1 = nn.BatchNorm2d(64)
+
+        if mnist:
+            self.conv1 = wrapped_conv(28, 1, 64, kernel_size=3, stride=1)
+            self.layer1 = self._make_layer(block, 28, 64, num_blocks[0], stride=1)
+            self.layer2 = self._make_layer(block, 28, 128, num_blocks[1], stride=2)
+            self.layer3 = self._make_layer(block, 14, 256, num_blocks[2], stride=2)
+            self.layer4 = self._make_layer(block, 7, 512, num_blocks[3], stride=2)
+        else:
+            self.conv1 = wrapped_conv(32, 3, 64, kernel_size=3, stride=1)
+            self.layer1 = self._make_layer(block, 32, 64, num_blocks[0], stride=1)
+            self.layer2 = self._make_layer(block, 32, 128, num_blocks[1], stride=2)
+            self.layer3 = self._make_layer(block, 16, 256, num_blocks[2], stride=2)
+            self.layer4 = self._make_layer(block, 8, 512, num_blocks[3], stride=2)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+        self.activation = F.leaky_relu if self.mod else F.relu
+        self.feature = None
+        self.temp = temp
+
+    def _make_layer(self, block, input_size, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(input_size, self.wrapped_conv, self.in_planes, planes, stride, self.mod,))
+            self.in_planes = planes * block.expansion
+            input_size = math.ceil(input_size / stride)
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        self.feature = out.clone().detach()
+        if self.temp==1:
+            out = self.fc(out)
+        else:
+            out = self.fc(out) / self.temp
+        return out
+
+
+def resnet18(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        BasicBlock,
+        [2, 2, 2, 2],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet50(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 6, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet101(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 23, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet110(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 26, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet152(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 8, 36, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model

SPC-UQ/Image_Classification/net/resnet_edl.py

@@ -0,0 +1,252 @@
+"""
+Pytorch implementation of ResNet models.
+Reference:
+[1] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR, 2016.
+"""
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+from net.spectral_normalization.spectral_norm_conv_inplace import spectral_norm_conv
+from net.spectral_normalization.spectral_norm_fc import spectral_norm_fc
+
+
+class AvgPoolShortCut(nn.Module):
+    def __init__(self, stride, out_c, in_c):
+        super(AvgPoolShortCut, self).__init__()
+        self.stride = stride
+        self.out_c = out_c
+        self.in_c = in_c
+
+    def forward(self, x):
+        if x.shape[2] % 2 != 0:
+            x = F.avg_pool2d(x, 1, self.stride)
+        else:
+            x = F.avg_pool2d(x, self.stride, self.stride)
+        pad = torch.zeros(x.shape[0], self.out_c - self.in_c, x.shape[2], x.shape[3], device=x.device,)
+        x = torch.cat((x, pad), dim=1)
+        return x
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, input_size, wrapped_conv, in_planes, planes, stride=1, mod=True):
+        super(BasicBlock, self).__init__()
+        self.conv1 = wrapped_conv(input_size, in_planes, planes, kernel_size=3, stride=stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = wrapped_conv(math.ceil(input_size / stride), planes, planes, kernel_size=3, stride=1)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.mod = mod
+        self.activation = F.leaky_relu if self.mod else F.relu
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            if mod:
+                self.shortcut = nn.Sequential(AvgPoolShortCut(stride, self.expansion * planes, in_planes))
+            else:
+                self.shortcut = nn.Sequential(
+                    wrapped_conv(input_size, in_planes, self.expansion * planes, kernel_size=1, stride=stride,),
+                    nn.BatchNorm2d(planes),
+                )
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = self.activation(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, input_size, wrapped_conv, in_planes, planes, stride=1, mod=True):
+        super(Bottleneck, self).__init__()
+        self.conv1 = wrapped_conv(input_size, in_planes, planes, kernel_size=1, stride=1)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = wrapped_conv(input_size, planes, planes, kernel_size=3, stride=stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = wrapped_conv(math.ceil(input_size / stride), planes, self.expansion * planes, kernel_size=1, stride=1)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+        self.mod = mod
+        self.activation = F.leaky_relu if self.mod else F.relu
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            if mod:
+                self.shortcut = nn.Sequential(AvgPoolShortCut(stride, self.expansion * planes, in_planes))
+            else:
+                self.shortcut = nn.Sequential(
+                    wrapped_conv(input_size, in_planes, self.expansion * planes, kernel_size=1, stride=stride,),
+                    nn.BatchNorm2d(self.expansion * planes),
+                )
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.activation(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = self.activation(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        block,
+        num_blocks,
+        num_classes=10,
+        temp=1.0,
+        spectral_normalization=True,
+        mod=True,
+        coeff=3,
+        n_power_iterations=1,
+        mnist=False,
+    ):
+        """
+        If the "mod" parameter is set to True, the architecture uses 2 modifications:
+        1. LeakyReLU instead of normal ReLU
+        2. Average Pooling on the residual connections.
+        """
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.mod = mod
+
+        def wrapped_conv(input_size, in_c, out_c, kernel_size, stride):
+            padding = 1 if kernel_size == 3 else 0
+
+            conv = nn.Conv2d(in_c, out_c, kernel_size, stride, padding, bias=False)
+
+            if not spectral_normalization:
+                return conv
+
+            # NOTE: Google uses the spectral_norm_fc in all cases
+            if kernel_size == 1:
+                # use spectral norm fc, because bound are tight for 1x1 convolutions
+                wrapped_conv = spectral_norm_fc(conv, coeff, n_power_iterations)
+            else:
+                # Otherwise use spectral norm conv, with loose bound
+                shapes = (in_c, input_size, input_size)
+                wrapped_conv = spectral_norm_conv(conv, coeff, shapes, n_power_iterations)
+
+            return wrapped_conv
+
+        self.wrapped_conv = wrapped_conv
+
+        self.bn1 = nn.BatchNorm2d(64)
+
+        if mnist:
+            self.conv1 = wrapped_conv(28, 1, 64, kernel_size=3, stride=1)
+            self.layer1 = self._make_layer(block, 28, 64, num_blocks[0], stride=1)
+            self.layer2 = self._make_layer(block, 28, 128, num_blocks[1], stride=2)
+            self.layer3 = self._make_layer(block, 14, 256, num_blocks[2], stride=2)
+            self.layer4 = self._make_layer(block, 7, 512, num_blocks[3], stride=2)
+        else:
+            self.conv1 = wrapped_conv(32, 3, 64, kernel_size=3, stride=1)
+            self.layer1 = self._make_layer(block, 32, 64, num_blocks[0], stride=1)
+            self.layer2 = self._make_layer(block, 32, 128, num_blocks[1], stride=2)
+            self.layer3 = self._make_layer(block, 16, 256, num_blocks[2], stride=2)
+            self.layer4 = self._make_layer(block, 8, 512, num_blocks[3], stride=2)
+        self.activation = F.leaky_relu if self.mod else F.relu
+        self.feature = None
+        self.temp = temp
+        self.output_alpha = nn.Linear(512 * block.expansion, num_classes)
+        self.output_beta = nn.Linear(512 * block.expansion, num_classes)
+        self.output_nu = nn.Linear(512 * block.expansion, num_classes)
+        self.output_gamma = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(self, block, input_size, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(input_size, self.wrapped_conv, self.in_planes, planes, stride, self.mod,))
+            self.in_planes = planes * block.expansion
+            input_size = math.ceil(input_size / stride)
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        self.feature = out.clone().detach()
+        alpha = self.output_alpha(out)
+        beta = self.output_beta(out)
+        nu = self.output_nu(out)
+        gamma = self.output_gamma(out)
+        alpha = F.softplus(alpha) + 1
+        beta = F.softplus(beta) + 1e-6
+        nu = F.softplus(nu) + 1e-6
+        gamma = F.softplus(gamma) + 1e-6
+        return alpha#[alpha, beta, nu, gamma]
+
+
+def resnet18_edl(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        BasicBlock,
+        [2, 2, 2, 2],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet50_edl(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 6, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet101_edl(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 23, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet110_edl(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 26, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet152_edl(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 8, 36, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model

SPC-UQ/Image_Classification/net/resnet_uq.py

@@ -0,0 +1,272 @@
+"""
+Pytorch implementation of ResNet models.
+Reference:
+[1] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR, 2016.
+"""
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+from net.spectral_normalization.spectral_norm_conv_inplace import spectral_norm_conv
+from net.spectral_normalization.spectral_norm_fc import spectral_norm_fc
+
+
+class AvgPoolShortCut(nn.Module):
+    def __init__(self, stride, out_c, in_c):
+        super(AvgPoolShortCut, self).__init__()
+        self.stride = stride
+        self.out_c = out_c
+        self.in_c = in_c
+
+    def forward(self, x):
+        if x.shape[2] % 2 != 0:
+            x = F.avg_pool2d(x, 1, self.stride)
+        else:
+            x = F.avg_pool2d(x, self.stride, self.stride)
+        pad = torch.zeros(x.shape[0], self.out_c - self.in_c, x.shape[2], x.shape[3], device=x.device,)
+        x = torch.cat((x, pad), dim=1)
+        return x
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, input_size, wrapped_conv, in_planes, planes, stride=1, mod=True):
+        super(BasicBlock, self).__init__()
+        self.conv1 = wrapped_conv(input_size, in_planes, planes, kernel_size=3, stride=stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = wrapped_conv(math.ceil(input_size / stride), planes, planes, kernel_size=3, stride=1)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.mod = mod
+        self.activation = F.leaky_relu if self.mod else F.relu
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            if mod:
+                self.shortcut = nn.Sequential(AvgPoolShortCut(stride, self.expansion * planes, in_planes))
+            else:
+                self.shortcut = nn.Sequential(
+                    wrapped_conv(input_size, in_planes, self.expansion * planes, kernel_size=1, stride=stride,),
+                    nn.BatchNorm2d(planes),
+                )
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = self.activation(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, input_size, wrapped_conv, in_planes, planes, stride=1, mod=True):
+        super(Bottleneck, self).__init__()
+        self.conv1 = wrapped_conv(input_size, in_planes, planes, kernel_size=1, stride=1)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = wrapped_conv(input_size, planes, planes, kernel_size=3, stride=stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = wrapped_conv(math.ceil(input_size / stride), planes, self.expansion * planes, kernel_size=1, stride=1)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+        self.mod = mod
+        self.activation = F.leaky_relu if self.mod else F.relu
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            if mod:
+                self.shortcut = nn.Sequential(AvgPoolShortCut(stride, self.expansion * planes, in_planes))
+            else:
+                self.shortcut = nn.Sequential(
+                    wrapped_conv(input_size, in_planes, self.expansion * planes, kernel_size=1, stride=stride,),
+                    nn.BatchNorm2d(self.expansion * planes),
+                )
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.activation(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = self.activation(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        block,
+        num_blocks,
+        num_classes=10,
+        temp=1.0,
+        spectral_normalization=True,
+        mod=True,
+        coeff=3,
+        n_power_iterations=1,
+        mnist=False,
+    ):
+        """
+        If the "mod" parameter is set to True, the architecture uses 2 modifications:
+        1. LeakyReLU instead of normal ReLU
+        2. Average Pooling on the residual connections.
+        """
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.mod = mod
+
+        def wrapped_conv(input_size, in_c, out_c, kernel_size, stride):
+            padding = 1 if kernel_size == 3 else 0
+
+            conv = nn.Conv2d(in_c, out_c, kernel_size, stride, padding, bias=False)
+
+            if not spectral_normalization:
+                return conv
+
+            # NOTE: Google uses the spectral_norm_fc in all cases
+            if kernel_size == 1:
+                # use spectral norm fc, because bound are tight for 1x1 convolutions
+                wrapped_conv = spectral_norm_fc(conv, coeff, n_power_iterations)
+            else:
+                # Otherwise use spectral norm conv, with loose bound
+                shapes = (in_c, input_size, input_size)
+                wrapped_conv = spectral_norm_conv(conv, coeff, shapes, n_power_iterations)
+
+            return wrapped_conv
+
+        self.wrapped_conv = wrapped_conv
+
+        self.bn1 = nn.BatchNorm2d(64)
+
+        if mnist:
+            self.conv1 = wrapped_conv(28, 1, 64, kernel_size=3, stride=1)
+            self.layer1 = self._make_layer(block, 28, 64, num_blocks[0], stride=1)
+            self.layer2 = self._make_layer(block, 28, 128, num_blocks[1], stride=2)
+            self.layer3 = self._make_layer(block, 14, 256, num_blocks[2], stride=2)
+            self.layer4 = self._make_layer(block, 7, 512, num_blocks[3], stride=2)
+        else:
+            self.conv1 = wrapped_conv(32, 3, 64, kernel_size=3, stride=1)
+            self.layer1 = self._make_layer(block, 32, 64, num_blocks[0], stride=1)
+            self.layer2 = self._make_layer(block, 32, 128, num_blocks[1], stride=2)
+            self.layer3 = self._make_layer(block, 16, 256, num_blocks[2], stride=2)
+            self.layer4 = self._make_layer(block, 8, 512, num_blocks[3], stride=2)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+        self.activation = F.leaky_relu if self.mod else F.relu
+        self.feature = None
+        self.temp = temp
+
+        def make_branch():
+            layers = []
+            in_features = 512 * block.expansion
+            neurons = 512 * block.expansion
+            for _ in range(1):
+                layers.append(nn.Linear(in_features, neurons))
+                layers.append(F.relu)
+                # layers.append(nn.Dropout(dropout_p))
+                in_features = neurons
+                # neurons //= 2
+            return nn.Sequential(*layers), nn.Linear(in_features, num_classes)
+
+        self.hidden_mar, self.mar = make_branch()
+        self.hidden_mar_up, self.mar_up = make_branch()
+        self.hidden_mar_down, self.mar_down = make_branch()
+
+
+    def forward(self, x):
+        mar = self.mar(self.hidden_mar(x))
+        mar_up = self.mar_up(self.hidden_mar_up(x))
+        mar_down = self.mar_down(self.hidden_mar_down(x))
+        return mar, mar_up, mar_down
+
+    def _make_layer(self, block, input_size, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(input_size, self.wrapped_conv, self.in_planes, planes, stride, self.mod,))
+            self.in_planes = planes * block.expansion
+            input_size = math.ceil(input_size / stride)
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = self.activation(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        self.feature = out.clone().detach()
+        if self.temp==1:
+            pred = self.fc(out)
+        else:
+            pred = self.fc(out) / self.temp
+
+        mar = self.mar(self.hidden_mar(out))
+        mar_up = self.mar_up(self.hidden_mar_up(out))
+        mar_down = self.mar_down(self.hidden_mar_down(out))
+        return pred, mar, mar_up, mar_down
+
+
+def resnet18(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        BasicBlock,
+        [2, 2, 2, 2],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet50(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 6, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet101(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 23, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet110(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 4, 26, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model
+
+
+def resnet152(spectral_normalization=True, mod=True, temp=1.0, mnist=False, imagenet=False, **kwargs):
+    model = ResNet(
+        Bottleneck,
+        [3, 8, 36, 3],
+        spectral_normalization=spectral_normalization,
+        mod=mod,
+        temp=temp,
+        mnist=mnist,
+        **kwargs
+    )
+    return model