Initial commit

README.md

@@ -0,0 +1,236 @@
+---
+license: gpl-3.0
+---
+
+# Model Card for ClotSimNet
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+This is the HuggingFace storage location for all of the models for the ClotSimNet project. They include several base models as well as their tuned variants. All models were trained in PyTorch.
+
+## Model Details
+
+### Model Description
+
+There are several models that were trained on the ClotSimNet dataset:
+
+* EfficientNet (B0, B3, B7)
+* ResNet (50, 152)
+* ConvNeXt-Large
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** Josh Gregory
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+Each model was imported using the instructions from the [PyTorch documentation](https://pytorch.org/vision/stable/models.html). The specific links for each model architecture are as follows:
+
+* [EfficientNet](https://proceedings.mlr.press/v97/tan19a.html)
+* [ResNet](https://ieeexplore.ieee.org/document/7780459)
+* [ConvNeXt-Large](https://ieeexplore.ieee.org/document/9879745)
+
+## Uses
+
+These models are intended to be used to predict permeabilities from porous simulations, such as those found in blood clots or other porous media. The training set was exclusively computational fluid dynamics (CFD) simulations of blood clots, however these models could be fine-tuned on other porous media datasets.
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+As of this publication, these models are not intended to treat, diagnose, or prevent any disease. They have not been approved by the FDA and are not intended for clinical or human use.
+
+## How to Get Started with the Model
+
+There are two flavors of each model. The base variants all had the same learning rate, weight decay, etc. Specifically for the CNNs:
+
+* Learning rate: 1e-3
+* Weight decay: 1e-5
+* Kernel size: 3
+* Stride: 2
+* Padding: 1
+* Use ImageNet pre-trained weights: True
+* Batch size: 2
+* Num workers: 70
+* Epochs: 500
+* Use image augmentation transforms: True
+* Learning rate reduction factor: 0.1
+* Learning rate reduction patience: 10
+* Learning rate reduction threshold: 1e-4
+
+For the MLPs, obviously things like the stride and padding are not applicable. Instead, the number of neurons per hidden layer and the number of hidden layers were set as:
+
+* Number of hidden layers: 3
+* Size of each hidden layer: 128
+
+Each model was exported in [ONNX](https://onnx.ai/) (`.onnx`) and [Safetensors](https://huggingface.co/docs/safetensors/index) (`.safetensors`). 
+
+### Loading the Models
+
+To import the ONNX (`.onnx`) file:
+
+```python
+import onnx
+import onnxruntime as ort
+import numpy as np
+
+
+# Load ONNX model
+model = onnx.load('path_to_model/model.onnx')
+
+# Validate model
+onnx.checker.check_model(onnx_model)
+
+# Create an inference session
+session = ort.InferenceSession('path_to_model/model.onnx')
+```
+
+TODO: Update inference to work with images
+
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

base/convnext_tiny/convnext_tiny_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5630e12bd9e1a749bb7582906632bc5729f39ff6c0c3778bf6d4a00cad1f21c
+size 111353138

base/convnext_tiny/convnext_tiny_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0551357284d36b50d964def3f8910ea047bbb80647baee087125d2b59b7abc9f
+size 111286724

base/enet_b0/enet_b0_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07515b6d3968deb03ff654ef7b4f7717e3cf1d8150eefc580c8ce41d471542a6
+size 16023020

base/enet_b0/enet_b0_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fe26ca7697ebf2064c1eced7866e31e1b6225c685d25b1e5b2156869eb318f3
+size 16238684

base/enet_b3/enet_b3_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23f5321a72e10a6f3ebdcf51137669b3d3f076da487c83faeb3cc5b604863b11
+size 42733583

base/enet_b3/enet_b3_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e68dc9870d1c3589e35390ae433d5c43eeebf9005b2399bdd7d6c0ee9019d590
+size 43197740

base/enet_b7/enet_b7_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20d0d19e6df3482854eda5c297f3a05aa49ff6c45e657ee84e09409e3d166654
+size 254785573

base/enet_b7/enet_b7_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6f46777f6d28b652ad1e83d3bbf8a33e15eadfdf6c18ad3bb6b5ef00eba9265
+size 256524508

base/mlp_dend/mlp_dend_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d13823b81798ffae5a4cae026a6ab65c07009279eb235eea6c3ed2b6597326e2
+size 76559

base/mlp_dend/mlp_dend_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:302e24376e0c459b4c7989eca39a09e52a3bd455c91eb50b35eb6db1e4b87cab
+size 76124

base/mlp_full/mlp_full_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70ae9dd8b401e4c5992390d4634a27c9e87efc27f9fb863646b66e07090f1c96
+size 203539

base/mlp_full/mlp_full_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:257ecba2a29380bb9996bcff2f293a666a418a6f9541010bbac6f5f16d5fcd2b
+size 203124

base/mlp_o1/mlp_o1_3d_pretrain_300_unfreeze_all.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e71ccd93151cd158dee6e41456cb8d21a893f367c4a3bd0927c460172197d0d1
+size 74511

base/mlp_o1/mlp_o1_3d_pretrain_300_unfreeze_all.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65c15a72b808418f92164a1c01550c9070bc25bac4f754e91afb7ad8b0065c0b
+size 74076

base/mlp_o1/mlp_o1_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5ec791f3288e7f620faac21424616f5d9f076ed8de0dcf7c6cde438dfed43d0
+size 74511

base/mlp_o1/mlp_o1_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efb4f86e381dc59bd0586674a7905a5124497ea26960a33190aed90b57be1bac
+size 74076

base/resnet_152/resnet_152_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26e83ab0b93023a40a23771999ef46cf10b0071c38e924f5a4949f3467986d06
+size 232317981

base/resnet_152/resnet_152_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cde5578dd31e177789eea1c073039ef566ca772a2117dc5b9dd768634f59be8
+size 233245860

base/resnet_18/resnet_18_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:caa4816f3191b19770ec71333a3cd1d694371429c6b412a212cea1416d3d27d1
+size 44663281

base/resnet_18/resnet_18_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94e3cc71115b126ff41213ea633196eeea0fe7799af5a1796388e56b66e9878c
+size 44723124

base/resnet_50/resnet_50_base.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:895e5cf66a159bd51e5f70a4c1f9bf4bebde9f69cfb572e2e0826f36dffba3bf
+size 93923245

base/resnet_50/resnet_50_base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad3f40f46be13f68448d6b87f4b6c7631fb886e187a16e3e7c62278ffe196b8e
+size 94248660

hyperparameters/best_params_convnext_tiny.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.23699388000283336,
+    "weight_decay": 0.0003795214141463416,
+    "kernel_size": 1,
+    "stride": 2,
+    "padding": 1,
+    "batch_size": 32,
+    "scheduler_factor": 0.10407032558978876,
+    "scheduler_patience": 14,
+    "scheduler_threshold": 0.0007639430825545345
+  },
+  "validation_loss": 0.36920997500419617
+}

hyperparameters/best_params_enet_b0.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.0015801849950581874,
+    "weight_decay": 1.235158048935322e-05,
+    "kernel_size": 1,
+    "stride": 2,
+    "padding": 0,
+    "batch_size": 64,
+    "scheduler_factor": 0.11559915010662485,
+    "scheduler_patience": 15,
+    "scheduler_threshold": 0.0003021977576971474
+  },
+  "validation_loss": 0.0038186332676559687
+}

hyperparameters/best_params_enet_b3.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.0011140616389973715,
+    "weight_decay": 7.02753409648415e-05,
+    "kernel_size": 5,
+    "stride": 1,
+    "padding": 2,
+    "batch_size": 32,
+    "scheduler_factor": 0.16322446851765238,
+    "scheduler_patience": 7,
+    "scheduler_threshold": 0.0038049686341583267
+  },
+  "validation_loss": 0.004008469637483358
+}

hyperparameters/best_params_enet_b7.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.0012199668475623272,
+    "weight_decay": 0.00032379826052113,
+    "kernel_size": 3,
+    "stride": 2,
+    "padding": 2,
+    "batch_size": 32,
+    "scheduler_factor": 0.11590671587549593,
+    "scheduler_patience": 10,
+    "scheduler_threshold": 0.0017160400686953745
+  },
+  "validation_loss": 0.002957067685201764
+}

hyperparameters/best_params_mlp.json

@@ -0,0 +1,13 @@
+{
+  "hyperparameters": {
+    "hidden_size": 1480,
+    "num_layers": 1,
+    "lr": 0.0008591530783482852,
+    "weight_decay": 0.00012686372039020111,
+    "batch_size": 32,
+    "scheduler_factor": 0.12387065250955309,
+    "scheduler_patience": 10,
+    "scheduler_threshold": 0.0001725777162252396
+  },
+  "validation_loss": 0.003614193992689252
+}

hyperparameters/best_params_mlp_dend.json

@@ -0,0 +1,13 @@
+{
+  "hyperparameters": {
+    "hidden_size": 1072,
+    "num_layers": 1,
+    "lr": 0.0006490624870220825,
+    "weight_decay": 0.0003579389623468425,
+    "batch_size": 32,
+    "scheduler_factor": 0.15137349637223427,
+    "scheduler_patience": 11,
+    "scheduler_threshold": 0.0002335875429607458
+  },
+  "validation_loss": 0.001131717232055962
+}

hyperparameters/best_params_mlp_o1.json

@@ -0,0 +1,13 @@
+{
+  "hyperparameters": {
+    "hidden_size": 1480,
+    "num_layers": 1,
+    "lr": 0.0008591530783482852,
+    "weight_decay": 0.00012686372039020111,
+    "batch_size": 32,
+    "scheduler_factor": 0.12387065250955309,
+    "scheduler_patience": 10,
+    "scheduler_threshold": 0.0001725777162252396
+  },
+  "validation_loss": 0.002175817731767893
+}

hyperparameters/best_params_resnet_152.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.0012199668475623272,
+    "weight_decay": 0.00032379826052113,
+    "kernel_size": 3,
+    "stride": 2,
+    "padding": 2,
+    "batch_size": 32,
+    "scheduler_factor": 0.11590671587549593,
+    "scheduler_patience": 10,
+    "scheduler_threshold": 0.0017160400686953745
+  },
+  "validation_loss": 0.0016866237856447697
+}

hyperparameters/best_params_resnet_18.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.0012199668475623272,
+    "weight_decay": 0.00032379826052113,
+    "kernel_size": 3,
+    "stride": 2,
+    "padding": 2,
+    "batch_size": 32,
+    "scheduler_factor": 0.11590671587549593,
+    "scheduler_patience": 10,
+    "scheduler_threshold": 0.0017160400686953745
+  },
+  "validation_loss": 0.002738025737926364
+}

hyperparameters/best_params_resnet_50.json

@@ -0,0 +1,14 @@
+{
+  "hyperparameters": {
+    "learning_rate": 0.0012199668475623272,
+    "weight_decay": 0.00032379826052113,
+    "kernel_size": 3,
+    "stride": 2,
+    "padding": 2,
+    "batch_size": 32,
+    "scheduler_factor": 0.11590671587549593,
+    "scheduler_patience": 10,
+    "scheduler_threshold": 0.0017160400686953745
+  },
+  "validation_loss": 0.002920560771599412
+}

model_classes/ConvNeXt.py

@@ -0,0 +1,528 @@
+import os
+from typing import Optional, List, Dict, Tuple, Any
+import numpy as np
+import pandas as pd
+import piexif
+from sklearn.preprocessing import StandardScaler
+from PIL import Image
+import torch
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, TensorDataset, DataLoader, random_split
+from torchvision.io import decode_image, read_file, image
+from torchvision.transforms import v2
+import torchvision.models as models
+import lightning as L
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+import torchmetrics
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch import Trainer, seed_everything
+import wandb
+import lmdb
+import msgpack
+import sys
+import torchvision
+
+class ConvNeXtLarge(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the ConvNeXt-Large architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.convnext_large(weights='DEFAULT')
+        else:
+            self.model = models.convnext_large()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.features[0][0] = nn.Conv2d(
+            in_channels=1, 
+            out_channels=192, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+    
+    # NOTE: pin_memory=True calls in the next three defs should be monitored for performance. Can easily set to false.
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+
+class ConvNeXtSmall(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the ConvNeXt-Small architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.convnext_small(weights='DEFAULT')
+        else:
+            self.model = models.convnext_small()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.features[0][0] = nn.Conv2d(
+            in_channels=1, 
+            out_channels=96, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+    
+    # NOTE: pin_memory=True calls in the next three defs should be monitored for performance. Can easily set to false.
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+    
+    
+class ConvNeXtTiny(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the ConvNeXt-Tiny architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.convnext_tiny(weights='DEFAULT')
+        else:
+            self.model = models.convnext_tiny()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.features[0][0] = nn.Conv2d(
+            in_channels=1, 
+            out_channels=96, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+    
+    # NOTE: pin_memory=True calls in the next three defs should be monitored for performance. Can easily set to false.
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)

model_classes/EfficientNets.py

@@ -0,0 +1,512 @@
+import os
+from typing import Optional, List, Dict, Tuple, Any
+import numpy as np
+import pandas as pd
+import piexif
+from sklearn.preprocessing import StandardScaler
+from PIL import Image
+import torch
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, TensorDataset, DataLoader, random_split
+from torchvision.io import decode_image, read_file, image
+from torchvision.transforms import v2
+import torchvision.models as models
+import lightning as L
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+import torchmetrics
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch import Trainer, seed_everything
+import wandb
+import lmdb
+import msgpack
+import sys
+import torchvision
+
+# Dynamically construct the path based on the user's home directory
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__name__), '..')))
+
+from utils import export_cnn, test_report
+
+torch.set_float32_matmul_precision('highest') # Options are medium, high, highest
+
+tensor_dtype = torch.bfloat16 # Can change to torch.bfloat16
+
+# For training on AMD GPUs, as bfloat16 isn't supported
+# tensor_dtype = torch.float16
+
+# EfficientNet B0 architecture
+class ENetB0(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the EfficientNet B0 architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+
+        if use_weights==True:
+            self.model = models.efficientnet_b0(weights='DEFAULT')
+        else:
+            self.model = models.efficientnet_b0()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.features[0][0] = nn.Conv2d(
+            in_channels=1, 
+            out_channels=32, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, 1)  # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+
+# EfficientNet B3 architecture
+class ENetB3(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the EfficientNet B3 architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.efficientnet_b3(weights='DEFAULT')
+        else:
+            self.model = models.efficientnet_b3()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.features[0][0] = nn.Conv2d(
+            in_channels=1, 
+            out_channels=40, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, 1)  # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'} 
+        
+# EfficientNet B7 architecture        
+class ENetB7(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4 
+    ):
+        '''
+        Class to load the EfficientNet B7 architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = True
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.efficientnet_b7(weights='DEFAULT')
+        else:
+            self.model = models.efficientnet_b7()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.features[0][0] = nn.Conv2d(
+            in_channels=1, 
+            out_channels=64, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.classifier[-1].in_features
+        self.model.classifier[-1] = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+  

model_classes/MLP.py

@@ -0,0 +1,186 @@
+import os
+from typing import Optional, List, Dict, Tuple, Any
+import numpy as np
+import pandas as pd
+import piexif
+from sklearn.preprocessing import StandardScaler
+from PIL import Image
+import torch
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, TensorDataset, DataLoader, random_split
+from torchvision.io import decode_image, read_file, image
+from torchvision.transforms import v2
+import torchvision.models as models
+import lightning as L
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+import torchmetrics
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch import Trainer, seed_everything
+import wandb
+import lmdb
+import msgpack
+import sys
+import torchvision
+
+# Dynamically construct the path based on the user's home directory
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__name__), '..')))
+
+from utils import export_mlp, test_report
+
+torch.set_float32_matmul_precision('highest') # Options are medium, high, highest
+
+tensor_dtype = torch.bfloat16 # Can change to torch.bfloat16
+
+# For training on AMD GPUs, as bfloat16 isn't supported
+# tensor_dtype = torch.float16
+
+class MLP(L.LightningModule):
+    def __init__(
+        self, 
+        input_size: int, 
+        hidden_size: int=128, 
+        num_layers: int=2,
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2,
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to create a simple multi-layered perceptron (MLP) model for tabular data.
+
+        Parameters:
+            input_size (int): Size of the input data.
+            hidden_size (int): Number of neurons for the hidden layers. Defaults to 100.
+            num_layers (int): Number of hidden layers. Defaults to 2.
+            learning_rate (float): Learning rate. Defaults to 1e-3
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-5
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters() # Save model hyperparameters for checkpointing
+
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.input_size = input_size
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+    
+        # Define model layers
+        layers = []
+        
+        layers.append(nn.Linear(self.input_size, hidden_size))
+        layers.append(nn.ReLU())
+        
+        # Set number of hidden layers dynamically via class definition
+        for _ in range(num_layers - 1):
+            layers.append(nn.Linear(hidden_size, hidden_size))
+            layers.append(nn.BatchNorm1d(num_features=hidden_size))
+            layers.append(nn.ReLU())
+            
+        layers.append(nn.Linear(hidden_size, 1)) # Output layer for regression and permeability
+        
+        self.model = nn.Sequential(*layers) # Build model
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+    def forward(self, x):
+        return self.model(x)
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}

model_classes/ResNets.py

@@ -0,0 +1,542 @@
+import os
+from typing import Optional, List, Dict, Tuple, Any
+import numpy as np
+import pandas as pd
+import piexif
+from sklearn.preprocessing import StandardScaler
+from PIL import Image
+import torch
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, TensorDataset, DataLoader, random_split
+from torchvision.io import decode_image, read_file, image
+from torchvision.transforms import v2
+import torchvision.models as models
+import lightning as L
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+import torchmetrics
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch import Trainer, seed_everything
+import wandb
+import lmdb
+import msgpack
+import sys
+import torchvision
+
+# Dynamically construct the path based on the user's home directory
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__name__), '..')))
+
+from utils import export_cnn, test_report
+
+torch.set_float32_matmul_precision('highest') # Options are medium, high, highest
+
+tensor_dtype = torch.bfloat16 # Can change to torch.bfloat16
+
+# For training on AMD GPUs, as bfloat16 isn't supported
+# tensor_dtype = torch.float16
+
+# ResNet-18 model architecture
+class ResNet18(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the ResNet-50 architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.resnet18(weights='DEFAULT')
+        else:
+            self.model = models.resnet18()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.conv1 = nn.Conv2d(
+            in_channels=1, 
+            out_channels=64, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.fc.in_features
+        self.model.fc = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+    
+    # NOTE: pin_memory=True calls in the next three defs should be monitored for performance. Can easily set to false.
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+# ResNet-50 model architecture
+class ResNet50(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the ResNet-50 architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.resnet50(weights='DEFAULT')
+        else:
+            self.model = models.resnet50()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.conv1 = nn.Conv2d(
+            in_channels=1, 
+            out_channels=64, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.fc.in_features
+        self.model.fc = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+    
+    # NOTE: pin_memory=True calls in the next three defs should be monitored for performance. Can easily set to false.
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+    
+# ResNet-152 architecture
+class ResNet152(L.LightningModule):
+    def __init__(
+        self, 
+        learning_rate: float=1e-3, 
+        weight_decay: float=1e-2, 
+        kernel_size: int=3, 
+        stride: int=1, 
+        padding: int=0,
+        use_weights: bool=True, 
+        scheduler_factor: float=0.1,
+        scheduler_patience: int=10,
+        scheduler_threshold: float=1e-4
+    ):
+        '''
+        Class to load the EfficientNet B0 architecture from PyTorch's torchvision library.
+
+        Parameters:
+            learning_rate (float): Learning rate. Defaults to 1e-3.
+            weight_decay (float): Weight decay for AdamW optimizer. Defaults to 1e-2.
+            kernel_size (int): Kernel size for convolutional filters. Defaults to 3.
+            stride (int): Stride for convolutional filters. Defaults to 1.
+            padding (int): Padding for convolutional filters. Defaults to 0.
+            use_weights (bool): Whether or not to use pretrained weights from ImageNet. Defaults to True.
+            scheduler_factor (float): Factor for ReduceLROnPlateau. Defaults to 0.1.
+            scheduler_patience (int): Patience for ReduceLROnPlateau. Defaults to 10.
+            scheduler_threshold (float): Threshold for ReduceLROnPlateau. Defaults to 1e-4.
+            
+        Methods:
+            All of the methods within this class are subclassed from PyTorch Lightning's LightningModule. Their documentation for each of these methods can be found here: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
+        '''
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Transforms flag to perform image augmentations
+        # self.use_transforms = use_transforms
+        
+        # Save predictions for later
+        self.test_preds = []
+        self.test_labels = []
+        
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.use_weights = use_weights
+        self.scheduler_factor = scheduler_factor
+        self.scheduler_patience = scheduler_patience
+        self.scheduler_threshold = scheduler_threshold
+        
+        if use_weights==True:
+            self.model = models.resnet152(weights='DEFAULT')
+        else:
+            self.model = models.resnet152()
+            
+        # Modify first conv. layer to accept grayscale/1 channel inputs
+        self.model.conv1 = nn.Conv2d(
+            in_channels=1, 
+            out_channels=64, 
+            kernel_size=kernel_size, 
+            stride=stride, 
+            padding=padding
+        )
+        
+        # Modify classifier for regression output
+        num_features = self.model.fc.in_features
+        self.model.fc = nn.Linear(num_features, 1) # Single output for regression
+        
+        # Create dictionary of metrics to track for training, validation, and testing
+        self.train_metrics = torchmetrics.MetricCollection(
+            {
+                'MAE': torchmetrics.regression.MeanAbsoluteError(num_outputs=1),
+                'RMSE': torchmetrics.regression.MeanSquaredError(squared=False, num_outputs=1),
+                'MSE': torchmetrics.regression.MeanSquaredError(squared=True, num_outputs=1),
+               'r_squared': torchmetrics.regression.R2Score() 
+            },
+            prefix='train_'
+        )
+        
+        self.valid_metrics = self.train_metrics.clone(prefix='valid_')
+        self.test_metrics = self.train_metrics.clone(prefix='test_')
+        
+        
+    def forward(self, x):
+        return self.model(x).squeeze(-1) # Remove extra dimension
+
+    # Reset metrics
+    def on_train_epoch_start(self):
+        self.train_metrics.reset()
+    
+    def training_step(self, batch, batch_idx):
+        x, y = batch # Get inputs and labels
+        y_pred = self(x) # Forward pass (call model's forward method)
+        
+        loss = F.mse_loss(y_pred, y)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True) # Log MSE loss for monitoring
+        
+        # Log metrics
+        batch_values = self.train_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        return loss
+
+    # Reset metrics
+    def on_validation_epoch_start(self):
+        self.valid_metrics.reset()
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.valid_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        return loss
+    
+    # Reset metrics
+    def on_test_epoch_start(self):
+        self.test_metrics.reset()
+        
+    def test_step(self, batch, batch_idx):
+        x, y = batch # Get features and labels
+        y_pred = self(x)
+        loss = F.mse_loss(y_pred, y)
+        self.log('test_loss', loss, on_step=True, on_epoch=True, logger=True, sync_dist=True)
+        
+        # Update validation metrics
+        batch_values = self.test_metrics(y_pred, y)
+        self.log_dict(batch_values, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # Store predictions and labels
+        self.test_preds.append(y_pred.cpu())
+        self.test_labels.append(y.cpu())
+        
+        return loss
+    
+    def on_test_epoch_end(self):
+        self.test_preds = torch.cat([pred.float() for pred in self.test_preds], dim=0).numpy()
+        self.test_labels = torch.cat([label.float() for label in self.test_labels], dim=0).numpy()
+
+        # Log transform test predictions and labels
+        self.test_preds = np.exp(self.test_preds)
+        self.test_labels = np.exp(self.test_labels)
+
+        # Store results for access after trainer.test()
+        self.test_results = {'preds': self.test_preds, 'labels': self.test_labels}
+        
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        scheduler = ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=self.scheduler_factor, patience=self.scheduler_patience, threshold=self.scheduler_threshold)
+        
+        return {'optimizer': optimizer,
+                'lr_scheduler': scheduler,
+                'monitor': 'val_loss'}
+    
+    # NOTE: pin_memory=True calls in the next three defs should be monitored for performance. Can easily set to false.
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)
+