Refactored and debugged model_runner.py

Author: joseph-hellerstein

requirements.txt

@@ -10,4 +10,5 @@ seaborn
 tabulate
 torch
 torchvision
+tqdm
 scikit-learn

src/iplane/autoencoder.py

@@ -10,30 +10,37 @@
 from typing import List
 import matplotlib.pyplot as plt
 
+LAYER_DIMENSIONS = [784, 512, 256, 128, 64]  # Example dimensions for MNIST
+
 
 ########################################################################
 class Autoencoder(nn.Module):
     # Basic Autoencoder
-    def __init__(self, input_dim=784, encoding_dim=64):
+    def __init__(self, layer_dimensions: List[int]):
+        """
+
+        Args:
+            dimensions (List[int]): List of dimensions for the autoencoder 
+                The first element is the input dimension,
+                the last element is the encoding dimension.
+        """
         super(Autoencoder, self).__init__()
+        self.layer_dimensions = layer_dimensions
+        self.input_dim = layer_dimensions[0]
+        self.encoding_dim = layer_dimensions[-1]
+        # Calculate dimension of hidden layer
         # Encoder
-        self.encoder = nn.Sequential(
-            nn.Linear(input_dim, 256),
-            nn.ReLU(),
-            nn.Linear(256, 128),
-            nn.ReLU(),
-            nn.Linear(128, encoding_dim),
-            nn.ReLU()
-        )
+        encoder_layers:list = []
+        for idx in range(len(layer_dimensions) - 1):
+            encoder_layers.append(nn.Linear(layer_dimensions[idx], layer_dimensions[idx + 1]))
+            encoder_layers.append(nn.ReLU())
+        self.encoder = nn.Sequential(*encoder_layers[0:-1]) 
         # Decoder
-        self.decoder = nn.Sequential(
-            nn.Linear(encoding_dim, 128),
-            nn.ReLU(),
-            nn.Linear(128, 256),
-            nn.ReLU(),
-            nn.Linear(256, input_dim),
-            nn.Sigmoid()  # Output between 0 and 1 for image reconstruction
-        )
+        decoder_layers:list = []
+        for idx in range(len(layer_dimensions) - 1, 0, -1):
+            decoder_layers.append(nn.Linear(layer_dimensions[idx], layer_dimensions[idx - 1]))
+            decoder_layers.append(nn.ReLU())
+        self.decoder = nn.Sequential(*decoder_layers[0:-1])
 
     def forward(self, x):
         # Encode
@@ -54,10 +61,12 @@ def decode(self, x):
 ########################################################################
 class AutoencoderRunner(object):
     # Runner for Autoencoder
+    layer_dimensions = [784, 512, 256, 128, 64]  # Example dimensions for MNIST
 
-    def __init__(self, num_epoch:int=3, learning_rate:float=1e-3, is_report:bool=False):
+    def __init__(self, layer_dimensions:List[int]=LAYER_DIMENSIONS,
+            num_epoch:int=3, learning_rate:float=1e-3, is_report:bool=False):
         self.device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"  # type: ignore
-        self.model = Autoencoder(input_dim=784, encoding_dim=64).to(self.device)
+        self.model = Autoencoder(layer_dimensions).to(self.device)
         self.num_epoch = num_epoch
         self.learning_rate = learning_rate
         self.losses: list = []

src/iplane/dataset_csv.py

@@ -1,13 +1,14 @@
 '''DataLoader for CSV files'''
 
+import numpy as np  # type: ignore
 import pandas as pd # type: ignore
 import torch
 from torch.utils.data import Dataset
-from typing import Optional, Union, cast
+from typing import Optional, Union
 
 
 class DatasetCSV(Dataset): 
-    def __init__(self, csv_input:Union[str, pd.DataFrame], target_column:str, transform=None):
+    def __init__(self, csv_input:Union[str, pd.DataFrame], target_column:Optional[str]=None, transform=None):
         """
         All columns except the target column are considered features.
 
@@ -20,9 +21,13 @@ def __init__(self, csv_input:Union[str, pd.DataFrame], target_column:str, transf
             self.data_df = csv_input
         else:
             self.data_df = pd.read_csv(csv_input)
+        self.target_column = target_column
         feature_columns = [col for col in self.data_df.columns if col != target_column]
         self.feature_tnsr = torch.tensor(self.data_df[feature_columns].values)
-        self.target_tnsr = torch.tensor(self.data_df[target_column].values)
+        if target_column is None:
+            self.target_tnsr = torch.zeros(len(self.data_df))
+        else:
+            self.target_tnsr = torch.tensor(self.data_df[target_column].values)
         self.transform = transform
 
     def __len__(self):
@@ -32,7 +37,10 @@ def __len__(self):
     def __getitem__(self, idx):
         # Get features and target
         feature_tnsr = self.feature_tnsr[idx].detach().clone()
-        target_tnsr = self.target_tnsr[idx].detach().clone()
+        if self.target_column is None:
+            target_tnsr = torch.tensor(np.nan)
+        else:
+            target_tnsr = self.target_tnsr[idx].detach().clone()
 
         # Apply transform if specified
         if self.transform:

src/iplane/model_runner.py

@@ -1,9 +1,11 @@
 
 from collections import namedtuple
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader
+from tqdm import tqdm  # type: ignore
 from typing import List, Tuple
 
 """To do
@@ -19,46 +21,78 @@ class ModelRunner(object):
     # Runner for Autoencoder
 
     def __init__(self, model: nn.Module, num_epoch:int=3, learning_rate:float=1e-3,
-                criterion:nn.Module=nn.MSELoss(), is_autoencoder:bool=False, is_report:bool=True):
+                criterion:nn.Module=nn.MSELoss(), is_autoencoder:bool=False,
+                is_normalized:bool=False, is_report:bool=True):
         """
         Args:
-            model (nn.Module): _description_
+            model (nn.Module): Model being run
             num_epoch (int, optional): Defaults to 3.
             learning_rate (float, optional): Defaults to 1e-3.
             is_autoencoder (bool, optional): target data is features Defaults to False.
-            is_report (bool, optional): Print text for progress. Defaults to False.
+            is_normalized (bool, optional): Whether to normalize the input data (divide by std).
+                                            Defaults to False.
+            is_report (bool, optional): Print text for progress.
+                                        Defaults to False.
         """
         self.device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"  # type: ignore
         self.model = model.to(self.device)
         self.num_epoch = num_epoch
         self.learning_rate = learning_rate
         self.criterion = criterion
         self.is_autoencoder = is_autoencoder
+        self.is_normalized = is_normalized
         self.is_report = is_report
+        # Calculated state
+        self.feature_std_tnsr = torch.tensor([np.nan])
+        self.target_std_tnsr = torch.tensor([np.nan])
 
     def train(self, train_loader: DataLoader) -> RunnerResult:
-        """Train the network."""
+        """
+        Train the model.
+
+        Args:
+            train_loader (DataLoader): DataLoader for training data
+        Returns:
+            RunnerResult: losses and number of epochs
+        """
+        ##
+        def calculate_std(loader_idx: int) -> torch.Tensor:
+            # loader_idx (int): Index into the DataLoader
+            full_tnsr = torch.cat([x[loader_idx] for x in train_loader])
+            if self.is_normalized:
+                return full_tnsr.std(dim=0)
+            else:
+                return torch.ones(full_tnsr.size()[1])
+        ##
+        # Handle normalization adjustments
+        self.feature_std_tnsr = calculate_std(0)
+        self.target_std_tnsr = calculate_std(1)
+        if self.is_autoencoder:
+            self.target_std_tnsr = self.feature_std_tnsr
+        # Initialize for training
         optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
         self.model.to(self.device)
-        
         self.model.train()
         losses = []
         avg_loss = 0.0
-        
-        for epoch in range(self.num_epoch):
+        epoch_loss = np.inf
+        # Training loop
+        pbar = tqdm(range(self.num_epoch), desc=f"epochs (loss={epoch_loss:.4f})")
+        for epoch in pbar:
+            pbar.set_description_str(f"epochs (loss={epoch_loss:.4f})")
             epoch_loss = 0
-            for data in list(train_loader):
-                data = data.to(self.device, non_blocking=True)
-                feature_tnsr, target_tnsr = data
+            #for idx, (feature_tnsr, target_tnsr) in list(train_loader):
+            for (feature_tnsr, target_tnsr) in train_loader:
                 if self.is_autoencoder:
                     # For autoencoder, target is the same as input
                     target_tnsr = feature_tnsr
-                feature_tnsr = feature_tnsr.view(feature_tnsr.size(0), -1)
-                # Forward pass
+                feature_tnsr = feature_tnsr/self.feature_std_tnsr
                 feature_tnsr = feature_tnsr.to(self.device)
+                target_tnsr = target_tnsr/self.target_std_tnsr
+                target_tnsr = target_tnsr.to(self.device)
+                # Forward pass
                 prediction_tnsr = self.model(feature_tnsr)
                 loss = self.criterion(prediction_tnsr, target_tnsr)
-                loss = loss.to(CPU)
                 # Backward pass
                 optimizer.zero_grad()
                 loss.backward()
@@ -71,22 +105,35 @@ def train(self, train_loader: DataLoader) -> RunnerResult:
             if self.is_report:
                 print(f'Epoch [{epoch+1}/{self.num_epoch}], Loss: {avg_loss:.4f}')
         #
+        self.model.to(CPU)
         return RunnerResult(losses=losses, num_epochs=self.num_epoch)
 
-    def evaluate(self, test_loader: DataLoader) -> RunnerResult:
-        """Evaluate the model on the test set."""
+    def predict(self, feature_tnsr: torch.Tensor) -> torch.Tensor:
+        """Predicts the target for the features.
+
+        Args:
+            feature_tnsr (torch.Tensor): Input features for which to predict targets.
+        Returns:
+            torch.Tensor: target predictions
+        """
+        self.model.eval()
+        feature_tnsr = feature_tnsr/self.feature_std_tnsr
+        with torch.no_grad():
+            prediction_tnsr = self.model(feature_tnsr)
+        return self.feature_std_tnsr*prediction_tnsr
+
+    def assess(self, test_loader: DataLoader) -> RunnerResult:
+        """Assess the model on a test dataset."""
         self.model.eval()
         test_losses = []
         #
         with torch.no_grad():
-            for data in list(test_loader):
-                data = data.to(self.device, non_blocking=True)
-                feature_tnsr, target_tnsr = data
+            for (feature_tnsr, target_tnsr) in list(test_loader):
                 if self.is_autoencoder:
                     # For autoencoder, target is the same as input
                     target_tnsr = feature_tnsr
                 feature_tnsr = feature_tnsr.view(feature_tnsr.size(0), -1)
-                prediction_tnsr = self.model(feature_tnsr)
+                prediction_tnsr = self.predict(feature_tnsr)
                 loss = self.criterion(prediction_tnsr, target_tnsr).to(CPU)
                 test_losses.append(loss.item())
 
@@ -110,5 +157,5 @@ def run(self, train_loader: DataLoader, test_loader: DataLoader)->Tuple[RunnerRe
             print("Training Fully Connected Autoencoder...")
         # Create and train fully connected autoencoder
         train_runner_result = self.train(train_loader)
-        test_runner_result = self.evaluate(test_loader)
+        test_runner_result = self.assess(test_loader)
         return train_runner_result, test_runner_result

tests/test_autoencoder.py

@@ -13,7 +13,7 @@
 NUM_EPOCH = 3
 
 
-class TestTrainingVisualizer(unittest.TestCase):
+class TestAutoencoder(unittest.TestCase):
 
     def setUp(self):
         self.runner = AutoencoderRunner(num_epoch=NUM_EPOCH, is_report=IGNORE_TESTS)