src.runner

  1import subprocess
  2
  3import numpy as np
  4import torch
  5from rich import print
  6from rich.panel import Panel
  7from torch import nn
  8from torch.nn import CrossEntropyLoss
  9from torch.optim import Adam
 10from torch.utils.data import DataLoader
 11from torch.utils.data import Dataset
 12from torch.utils.tensorboard.writer import SummaryWriter
 13
 14from src.data import get_mnist_dataloaders
 15from src.eval import eval_iteration
 16from src.model import BasicCNN
 17from src.plot import generate_projection
 18from src.plot import plot_confusion_matrix
 19from src.plot import saliency_map
 20from src.train import generate_canocial
 21from src.train import training_iteration
 22
 23# set global seed
 24torch.manual_seed(0)
 25
 26
 27def run_pipeline(
 28    model: nn.Module,
 29    dataloader_train: DataLoader,
 30    dataloader_val: DataLoader,
 31    dataset_train: Dataset,
 32    dataset_val: Dataset,
 33    n_epochs: int,
 34    min_loss: float,
 35    writer: SummaryWriter,
 36    model_weights_path: str,
 37):
 38    """The full training pipeline performs a loop over the following:
 39
 40    ```Loop
 41    1) Train the model
 42    2) Perform evaluation
 43    3) Record the average loss
 44        If the model has the best average loss:
 45            Save off the model
 46            Generate a confusion matrix
 47        Else:
 48            continue training
 49    ```
 50
 51    .. note::
 52        The `dataloaders` must provide `(x, y)` representing inputs and true labels.
 53
 54    Args:
 55        model (nn.Module): Torch model.
 56        dataloader_train (DataLoader): DataLoader with `(x, y)` data.
 57        dataloader_val (DataLoader): DataLoader with `(x, y)` data.
 58        dataset_train (Dataset): Datset with `(x, y)` data.
 59        dataset_val (Dataset): Datset with `(x, y)` data.
 60        n_epochs (int): Number of epochs to run training.
 61        min_loss (float): Minimium loss recorded for the training iterations.
 62        writer (SummaryWriter): Tensorboard writer.
 63        model_weights_path (str): Path to save off model.
 64    """
 65    # Create optimizer and loss functions
 66    optimizer = Adam(params=model.parameters(), lr=0.00001)
 67    loss_fn = CrossEntropyLoss()
 68
 69    # Perform training on the data
 70    for global_step in range(0, n_epochs):
 71        print(f"Global step: {global_step}")
 72
 73        # Perform training
 74        avg_train_loss = training_iteration(
 75            model=model,
 76            loss_fn=loss_fn,
 77            optimizer=optimizer,
 78            dataloader_train=dataloader_train,
 79            global_step=global_step,
 80            writer=writer,
 81        )
 82
 83        # Perform evaluation
 84        avg_eval_loss = eval_iteration(
 85            model=model,
 86            loss_fn=loss_fn,
 87            dataloader_val=dataloader_val,
 88            global_step=global_step,
 89            writer=writer,
 90        )
 91
 92        # Printout
 93        print(f"Epoch:               {global_step}/{n_epochs-1}")
 94        print(f"    Training   Loss: {avg_train_loss:.3f}")
 95        print(f"    Evaluation Loss: {avg_eval_loss:.3f}")
 96
 97        # Save off best model
 98        if avg_eval_loss < min_loss:
 99            # Update the min_loss as the new baseline
100            min_loss = avg_eval_loss
101
102            # Save off the best model
103            torch.save(model.state_dict(), model_weights_path)
104
105            # Run analysis
106            run_analysis(
107                model=model,
108                dataloader_val=dataloader_val,
109                dataset_val=dataset_val,
110                global_step=global_step,
111                writer=writer,
112            )
113
114            # Printout
115            print(f"Saving best model with evaluation loss of: {avg_eval_loss:.2f}")
116
117
118def run_analysis(
119    model: nn.Module,
120    dataloader_val: DataLoader,
121    dataset_val: Dataset,
122    global_step: int,
123    writer: SummaryWriter,
124):
125    """Run analysis on each epoch where the model improves.  This should log multiple metrics that can be used in post-analysis examination.
126
127    Args:
128        model (nn.Module): Torch model.
129        dataloader_val (DataLoader): Dataloader for validation.
130        dataset_val (Dataset): Datset for validation.
131        global_step (int): Global step for tracking.
132        writer (SummaryWriter): Tensorboard writer.
133    """
134
135    # Save off confusion matrix
136    plot_confusion_matrix(
137        model=model,
138        dataloader_val=dataloader_val,
139        save_path="save/plots/confusion_matrix",
140        fig_name=f"confusion_matrix_{global_step}",
141    )
142
143    # Generate projection logs
144    generate_projection(
145        model=model,
146        dataloader_val=dataloader_val,
147        writer=writer,
148        global_step=global_step,
149    )
150
151    # Generate saliency plot
152    saliency_map(
153        model=model,
154        dataset_val=dataset_val,
155        figure_name="saliency",
156        global_step=global_step,
157        save_path=f"save/plots/saliency",
158    )
159
160
161def run_mnist(
162    n_epochs: int = 2,
163    min_loss: float = np.inf,
164    batch_size: int = 32,
165    path_model: str = "save/models",
166    path_log: str = "save/logs",
167    path_plots: str = "save/plots",
168    model_name: str = "best_model.pkl",
169):
170    """Perform MNIST training.  This function performs the following:
171
172        1) Create save/log directories
173        2) Get data iterators for train and evaluation
174        3) Create Tensorboard writer
175        4) Create model
176        5) Execute training pipeline using the data iterators, Tensorboard writer, and model provided
177
178    Args:
179        n_epochs (int, optional): Number of epochs to perform training. Defaults to 2.
180        min_loss (float, optional): The minimum loss recorded thus far. Defaults to np.inf.
181        batch_size (int, optional): The batch size to use when retrieving data from data iterators. Defaults to 16.
182        path_model (str, optional): The path to the model. Defaults to "save/models".
183        path_log (str, optional): The path to the logs. Defaults to "logs".
184        model_name (str, optional): The name of the model. Defaults to "best_model.pkl".
185    """
186    # Fully defined path to model save file
187    model_weights_path = f"{path_model}/{model_name}"
188
189    # Create directories
190    subprocess.run(["mkdir", "-p", path_model])
191    subprocess.run(["mkdir", "-p", path_log])
192    subprocess.run(["mkdir", "-p", path_plots])
193
194    # Preparing pipeline
195    dataloader_train, dataloader_val, dataset_train, dataset_val = get_mnist_dataloaders(
196        batch_size=batch_size
197    )
198    writer = SummaryWriter(path_log)  # Tensorboard
199    model = BasicCNN()  # hyperparameters
200
201    # Preprocessing
202    saliency_map(
203        model,
204        dataset_val=dataset_val,
205        figure_name="saliency",
206        global_step=-1,
207        save_path=f"{path_plots}/saliency",
208    )
209
210    # Save off model architecture.  Requires that you pass in a torch tensor for forward pass in order for logging to happen.
211    writer.add_graph(model, torch.rand(1, 1, 28, 28))
212
213    # Train the model
214    run_pipeline(
215        model=model,
216        dataloader_train=dataloader_train,
217        dataloader_val=dataloader_val,
218        dataset_train=dataset_train,
219        dataset_val=dataset_val,
220        n_epochs=n_epochs,
221        min_loss=min_loss,
222        writer=writer,
223        model_weights_path=model_weights_path,
224    )
225
226    # TODO: generate_canocial is not currently working
227    writer.close()
228
229    # Printout
230    print(Panel("To open Tensorboard type:"))
231    print("     tensorboard --logdir=save/logs")
232    print("     :warning: Chrome is recommended for viewing projections!")
233
234
235if __name__ == "__main__":
236    run_mnist()