src.plot

  1import subprocess
  2from copy import deepcopy
  3
  4import matplotlib.pyplot as plt
  5import torch
  6from sklearn.metrics import confusion_matrix
  7from sklearn.metrics import ConfusionMatrixDisplay
  8from torch import nn
  9from torch.utils.data import DataLoader
 10from torch.utils.data import Dataset
 11from torch.utils.tensorboard.writer import SummaryWriter
 12
 13
 14def plot_confusion_matrix(
 15    model: nn.Module,
 16    dataloader_val: DataLoader,
 17    save_path: str = "save",
 18    fig_name: str = "",
 19    fig_x_size: int = 8,
 20    fig_y_size: int = 8,
 21    dpi: int = 250,
 22):
 23    """Given a model and a DataLoader (which provides inputs and labels), evaluate how well the model makes predictions on the data.  We compare the highest predicted value against the true label and generate a confusion matrix.
 24
 25    Args:
 26        model (nn.Module): Torch model.
 27        dataloader_val (DataLoader): DataLoader that outputs `(x, y)`.
 28        save_path (str, optional): Path to save plots. Defaults to "save".
 29        fig_name (str, optional): Figure name. Defaults to "".
 30        fig_x_size (int, optional): Figure x size. Defaults to 8.
 31        fig_y_size (int, optional): Figure y size. Defaults to 8.
 32        dpi (int, optional): Increase for higher resolution. Defaults to 250.
 33    """
 34
 35    # Make dirs
 36    subprocess.run(["mkdir", "-p", save_path])
 37
 38    # Preallocate lists
 39    y_pred = []
 40    y_true = []
 41
 42    # Iterate over all inputs and labels for the data iterator
 43    for x, y in dataloader_val:
 44        # Perform a forward pass with the input data x and get highest predicted value
 45        y_hat, _ = model(x)
 46        y_idx = torch.argmax(y_hat, 1)
 47
 48        # Add the prediction and true label to the preallocated list
 49        y_pred.extend(list(y_idx.numpy()))
 50        y_true.extend(list(y.numpy()))
 51
 52    # Generate the confusion matrix using the predictions versus the true labels
 53    cm = confusion_matrix(y_true, y_pred)
 54
 55    # Generate a figure and plot the results
 56    figure = plt.figure(figsize=(fig_x_size, fig_y_size))
 57    disp = ConfusionMatrixDisplay(confusion_matrix=cm)
 58    disp.plot()
 59
 60    # Add axes labels
 61    plt.title("Confusion Matrix")
 62    plt.ylabel("Predictions")
 63    plt.xlabel("Truth")
 64
 65    # Save off the plot
 66    plt.savefig(f"{save_path}/{fig_name}.png", dpi=dpi)
 67    plt.close("all")
 68
 69
 70def generate_projection(
 71    model: nn.Module,
 72    dataloader_val: DataLoader,
 73    writer: SummaryWriter,
 74    global_step: int = 0,
 75    projection_limit: int = 1000,
 76):
 77    """Create a projection of the data in N-dimensional space.  This can be visualized using Tensorboard.
 78
 79    Args:
 80        model (nn.Module): Torch model.
 81        dataloader_val (DataLoader): Dataloader that outputs `(x, y)`.
 82        writer (SummaryWriter): Tensorboard writer.
 83        global_step (int, optional): The global step tracker. Defaults to 0.
 84        projection_limit (int, optional): The maximum number of projections allowed. Defaults to 1000.
 85    """
 86    # Initialize variables
 87    embeddings = []
 88    labels = []
 89    images = []
 90
 91    # Iterate over all data
 92    for x, y in dataloader_val:
 93        # Remove gradient tape
 94        with torch.no_grad():
 95            _, embedding = model(x)  # get the embeddings (not predictions of the model)
 96            embeddings.append(embedding)  # save off embeddings
 97
 98        # Append the data
 99        labels.append(y)
100        images.append(x)
101
102    # Stack the output embeddings
103    stacked_embeddings = torch.vstack(embeddings)
104    stacked_labels = torch.hstack(labels).numpy()  # convert to numpy array instead of tensors
105    stacked_images = torch.vstack(images)
106
107    # Tensorboard is only able to plot a maximum number of data points in its projection.  Thus, we will limit the number of data points we will display.
108    mean, std, _ = (
109        torch.mean(stacked_embeddings),
110        torch.std(stacked_embeddings),
111        torch.var(stacked_embeddings),
112    )
113    stacked_embeddings = (stacked_embeddings - mean) / std
114    stacked_embeddings = stacked_embeddings[0:projection_limit, :]
115    stacked_labels = stacked_labels[0:projection_limit]
116    stacked_images = stacked_images[0:projection_limit, :, :, :]
117
118    # Add the embedding information to Tensorboard logs
119    writer.add_embedding(
120        stacked_embeddings,
121        metadata=stacked_labels,
122        # label_img=stacked_images,     # turn on to display images for each data point!
123        global_step=global_step,
124    )
125
126
127def saliency_map(
128    model: nn.Module,
129    dataset_val: Dataset,
130    figure_name: str = "saliency",
131    global_step: int = 0,
132    model_weights_path: str | None = None,
133    save_path: str = "save",
134):
135    """Create a saliency map showing which inputs contributed the most toward the predictions.
136
137    Args:
138        model (nn.Module): Torch model.
139        dataset_val (Dataset): The Dataset which provides `(x, y)`.
140        figure_name (str): The name of the figure to save as file. Defaults to "saliency".
141        global_step (int): The global step for tracking. Defaults to 0.
142        model_weights_path (str | None): Torch model weights path.
143        save_path (str): Save path.
144    """
145
146    # Create dirs if they do not exists
147    subprocess.run(["mkdir", "-p", save_path])
148
149    # configuration (must be even number!)
150    n_rows = 5
151    n_cols = 10
152
153    # Copy model
154    nn_model = deepcopy(model)
155    if model_weights_path is not None:  # Load weights into model if available
156        nn_model.load_state_dict(torch.load(model_weights_path))
157
158    data_iter = iter(dataset_val)
159    data_idx = 0
160
161    for ii in range(0, n_rows, 5):
162        for jj in range(1, n_cols + 1, 1):
163            x, y = next(data_iter)
164            data_idx += 1
165
166            # Calculate gradients w.r.t. input from output
167            x_base = torch.ones((1, 28, 28)) * -1
168            x_pred = x
169
170            x_base.requires_grad = True  # set gradient tape to True
171            x_pred.requires_grad = True  # set gradient tape to True
172
173            y_base, _ = nn_model(x_base)  # forward prop baseline
174            y_pred, _ = nn_model(x_pred)  # forward prop
175
176            y_base.sum().backward()  # baseline backpropagation
177            y_pred.sum().backward()  # prediction backpropagation
178
179            # Normalize the saliency plot
180            img_source = x.squeeze().detach().numpy()
181            img_baseline = torch.abs(x_base.grad.squeeze())
182            img_saliency = torch.abs(x_pred.grad.squeeze())
183            img_delta = x_pred.grad.squeeze() - x_base.grad.squeeze()
184            img_overlay = img_saliency * img_source
185
186            # Generate subplots
187            column_idx = ii * n_cols + jj
188            ax1 = plt.subplot(n_rows, n_cols, column_idx + 0 * n_cols)
189            ax2 = plt.subplot(n_rows, n_cols, column_idx + 1 * n_cols)
190            ax3 = plt.subplot(n_rows, n_cols, column_idx + 2 * n_cols)
191            ax4 = plt.subplot(n_rows, n_cols, column_idx + 3 * n_cols)
192            ax5 = plt.subplot(n_rows, n_cols, column_idx + 4 * n_cols)
193
194            # plot images
195            ax1.imshow(img_source, cmap=plt.cm.viridis, aspect="auto")
196            ax2.imshow(img_baseline, cmap=plt.cm.viridis, aspect="auto")
197            ax3.imshow(img_saliency, cmap=plt.cm.viridis, aspect="auto")
198            ax4.imshow(img_delta, cmap=plt.cm.viridis, aspect="auto")
199            ax5.imshow(img_overlay, cmap=plt.cm.viridis, aspect="auto")
200
201            # Add labels for first column only
202            if column_idx == 1:
203                ax1.set_ylabel("source")
204                ax2.set_ylabel("baseline")
205                ax3.set_ylabel("saliency")
206                ax4.set_ylabel("delta")
207                ax5.set_ylabel("overlay")
208
209            # Remove all ticks
210            remove_ticks(ax1)
211            remove_ticks(ax2)
212            remove_ticks(ax3)
213            remove_ticks(ax4)
214            remove_ticks(ax5)
215
216    # Save figure
217    plt.savefig(f"{save_path}/{figure_name}_{global_step}.png", dpi=250)
218
219
220from typing import Type
221
222
223def remove_ticks(axis):
224    axis.tick_params(
225        axis="y",  # changes apply to the x-axis
226        which="both",  # both major and minor ticks are affected
227        left=False,  # ticks along the bottom edge are off
228        labelleft=False,  # labels along the left edge are off
229    )
230
231    axis.tick_params(
232        axis="x",  # changes apply to the x-axis
233        which="both",  # both major and minor ticks are affected
234        bottom=False,  # ticks along the bottom edge are off
235        top=False,  # ticks along the top edge are off
236        labelbottom=False,  # labels along the bottom edge are off
237    )