Spaces:
Runtime error
Runtime error
import matplotlib.pyplot as plt | |
import numpy as np | |
import pandas as pd | |
import pytorch_lightning as pl | |
import seaborn as sns | |
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
import torch.optim as optim | |
import torchmetrics | |
from torch.optim.lr_scheduler import OneCycleLR | |
from torch_lr_finder import LRFinder | |
from . import config | |
from .visualize import plot_incorrect_preds | |
class Net(pl.LightningModule): | |
def __init__( | |
self, | |
num_classes=10, | |
dropout_percentage=0, | |
norm="bn", | |
num_groups=2, | |
criterion=F.cross_entropy, | |
learning_rate=0.001, | |
weight_decay=0.0, | |
): | |
super(Net, self).__init__() | |
if norm == "bn": | |
self.norm = nn.BatchNorm2d | |
elif norm == "gn": | |
self.norm = lambda in_dim: nn.GroupNorm( | |
num_groups=num_groups, num_channels=in_dim | |
) | |
elif norm == "ln": | |
self.norm = lambda in_dim: nn.GroupNorm(num_groups=1, num_channels=in_dim) | |
# Define the loss criterion | |
self.criterion = criterion | |
# Define the Metrics | |
self.accuracy = torchmetrics.Accuracy( | |
task="multiclass", num_classes=num_classes | |
) | |
self.confusion_matrix = torchmetrics.ConfusionMatrix( | |
task="multiclass", num_classes=config.NUM_CLASSES | |
) | |
# Define the Optimizer Hyperparameters | |
self.learning_rate = learning_rate | |
self.weight_decay = weight_decay | |
# Prediction Storage | |
self.pred_store = { | |
"test_preds": torch.tensor([]), | |
"test_labels": torch.tensor([]), | |
"test_incorrect": [], | |
} | |
self.log_store = { | |
"train_loss_epoch": [], | |
"train_acc_epoch": [], | |
"val_loss_epoch": [], | |
"val_acc_epoch": [], | |
"test_loss_epoch": [], | |
"test_acc_epoch": [], | |
} | |
# This defines the structure of the NN. | |
# Prep Layer | |
self.prep_layer = nn.Sequential( | |
nn.Conv2d(3, 64, kernel_size=3, padding=1), # 32x32x3 | 1 -> 32x32x64 | 3 | |
self.norm(64), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
) | |
self.l1 = nn.Sequential( | |
nn.Conv2d(64, 128, kernel_size=3, padding=1), # 32x32x128 | 5 | |
nn.MaxPool2d(2, 2), # 16x16x128 | 6 | |
self.norm(128), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
) | |
self.l1res = nn.Sequential( | |
nn.Conv2d(128, 128, kernel_size=3, padding=1), # 16x16x128 | 10 | |
self.norm(128), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
nn.Conv2d(128, 128, kernel_size=3, padding=1), # 16x16x128 | 14 | |
self.norm(128), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
) | |
self.l2 = nn.Sequential( | |
nn.Conv2d(128, 256, kernel_size=3, padding=1), # 16x16x256 | 18 | |
nn.MaxPool2d(2, 2), # 8x8x256 | 19 | |
self.norm(256), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
) | |
self.l3 = nn.Sequential( | |
nn.Conv2d(256, 512, kernel_size=3, padding=1), # 8x8x512 | 27 | |
nn.MaxPool2d(2, 2), # 4x4x512 | 28 | |
self.norm(512), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
) | |
self.l3res = nn.Sequential( | |
nn.Conv2d(512, 512, kernel_size=3, padding=1), # 4x4x512 | 36 | |
self.norm(512), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
nn.Conv2d(512, 512, kernel_size=3, padding=1), # 4x4x512 | 44 | |
self.norm(512), | |
nn.ReLU(), | |
nn.Dropout(dropout_percentage), | |
) | |
self.maxpool = nn.MaxPool2d(4, 4) | |
# Classifier | |
self.linear = nn.Linear(512, 10) | |
def forward(self, x): | |
x = self.prep_layer(x) | |
x = self.l1(x) | |
x = x + self.l1res(x) | |
x = self.l2(x) | |
x = self.l3(x) | |
x = x + self.l3res(x) | |
x = self.maxpool(x) | |
x = x.view(-1, 512) | |
x = self.linear(x) | |
return F.log_softmax(x, dim=1) | |
def training_step(self, batch, batch_idx): | |
data, target = batch | |
# print("curr lr: ", self.optimizers().param_groups[0]["lr"]) | |
# forward pass | |
pred = self(data) | |
# Calculate loss | |
loss = self.criterion(pred, target) | |
# Calculate the metrics | |
accuracy = self.accuracy(pred, target) | |
self.log_dict( | |
{"train_loss": loss, "train_acc": accuracy}, | |
on_step=True, | |
on_epoch=True, | |
prog_bar=True, | |
logger=True, | |
) | |
return loss | |
def validation_step(self, batch, batch_idx): | |
data, target = batch | |
# forward pass | |
pred = self(data) | |
# Calculate loss | |
loss = self.criterion(pred, target) | |
# Calculate the metrics | |
accuracy = self.accuracy(pred, target) | |
self.log_dict( | |
{"val_loss": loss, "val_acc": accuracy}, | |
on_step=True, | |
on_epoch=True, | |
prog_bar=True, | |
logger=True, | |
) | |
return loss | |
def test_step(self, batch, batch_idx): | |
data, target = batch | |
# forward pass | |
pred = self(data) | |
argmax_pred = pred.argmax(dim=1).cpu() | |
# Calculate loss | |
loss = self.criterion(pred, target) | |
# Calculate the metrics | |
accuracy = self.accuracy(pred, target) | |
self.log_dict( | |
{"test_loss": loss, "test_acc": accuracy}, | |
on_step=True, | |
on_epoch=True, | |
prog_bar=True, | |
logger=True, | |
) | |
# Update the confusion matrix | |
self.confusion_matrix.update(pred, target) | |
# Store the predictions, labels and incorrect predictions | |
data, target, pred, argmax_pred = ( | |
data.cpu(), | |
target.cpu(), | |
pred.cpu(), | |
argmax_pred.cpu(), | |
) | |
self.pred_store["test_preds"] = torch.cat( | |
(self.pred_store["test_preds"], argmax_pred), dim=0 | |
) | |
self.pred_store["test_labels"] = torch.cat( | |
(self.pred_store["test_labels"], target), dim=0 | |
) | |
for d, t, p, o in zip(data, target, argmax_pred, pred): | |
if p.eq(t.view_as(p)).item() == False: | |
self.pred_store["test_incorrect"].append( | |
(d.cpu(), t, p, o[p.item()].cpu()) | |
) | |
return loss | |
def find_bestLR_LRFinder(self, optimizer): | |
lr_finder = LRFinder(self, optimizer, criterion=self.criterion) | |
lr_finder.range_test( | |
self.trainer.datamodule.train_dataloader(), | |
end_lr=config.LRFINDER_END_LR, | |
num_iter=config.LRFINDER_NUM_ITERATIONS, | |
step_mode=config.LRFINDER_STEP_MODE, | |
) | |
best_lr = None | |
try: | |
_, best_lr = lr_finder.plot() # to inspect the loss-learning rate graph | |
except Exception as e: | |
pass | |
lr_finder.reset() # to reset the model and optimizer to their initial state | |
return best_lr | |
def configure_optimizers(self): | |
optimizer = self.get_only_optimizer() | |
best_lr = self.find_bestLR_LRFinder(optimizer) | |
scheduler = OneCycleLR( | |
optimizer, | |
max_lr=1.47e-03, | |
# total_steps=self.trainer.estimated_stepping_batches, | |
steps_per_epoch=len(self.trainer.datamodule.train_dataloader()), | |
epochs=config.NUM_EPOCHS, | |
pct_start=5 / config.NUM_EPOCHS, | |
div_factor=config.OCLR_DIV_FACTOR, | |
three_phase=config.OCLR_THREE_PHASE, | |
final_div_factor=config.OCLR_FINAL_DIV_FACTOR, | |
anneal_strategy=config.OCLR_ANNEAL_STRATEGY, | |
) | |
return [optimizer], [ | |
{"scheduler": scheduler, "interval": "step", "frequency": 1} | |
] | |
def get_only_optimizer(self): | |
optimizer = optim.Adam( | |
self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay | |
) | |
return optimizer | |
def on_test_end(self) -> None: | |
super().on_test_end() | |
## Confusion Matrix | |
confmat = self.confusion_matrix.cpu().compute().numpy() | |
if config.NORM_CONF_MAT: | |
df_confmat = pd.DataFrame( | |
confmat / np.sum(confmat, axis=1)[:, None], | |
index=[i for i in config.CLASSES], | |
columns=[i for i in config.CLASSES], | |
) | |
else: | |
df_confmat = pd.DataFrame( | |
confmat, | |
index=[i for i in config.CLASSES], | |
columns=[i for i in config.CLASSES], | |
) | |
plt.figure(figsize=(7, 5)) | |
sns.heatmap(df_confmat, annot=True, cmap="Blues", fmt=".3f", linewidths=0.5) | |
plt.tight_layout() | |
plt.ylabel("True label") | |
plt.xlabel("Predicted label") | |
plt.show() | |
def plot_incorrect_predictions_helper(self, num_imgs=10): | |
return plot_incorrect_preds( | |
self.pred_store["test_incorrect"], config.CLASSES, num_imgs | |
) | |