Some updates in code after workshop

Author: tfha

src/2022_06_10_analyze_rock_classification.py

@@ -1,4 +1,9 @@
-"""Script for analysis of rock type classification model"""
+"""Script for analysis of rock type classification model
+
+Run tensorboard calling this in your terminal:
+tensorboard --logdir Reports/tensorboard_logs
+
+"""
 
 import torch
 from pathlib import Path
@@ -21,8 +26,9 @@
 disp = ConfusionMatrixDisplay(confusion_matrix=cm)
 fig, ax = plt.subplots(figsize=(15,15))
 disp.plot(cmap="viridis", ax=ax, values_format=".2f")
+# ax.set_xticklabels(perf["class_names"])
 ax.set_xlabel("Predicted rocktype")
 ax.set_ylabel("True rocktype")
 plt.tight_layout()
-plt.savefig("Figures/confusion_matrix_rocktypes.png", dpi=600)
-# plt.show()
+plt.savefig(Path("Figures/confusion_matrix_rocktypes.png"), dpi=600)
+plt.show()

src/2022_06_10_image_classification.py

@@ -8,13 +8,17 @@
 https://www.kaggle.com/datasets/neelgajare/rocks-dataset
 
 TODO: Try to classify better with:
-- better tuning of parameters
-- weighting in loss function
+- Tuning of hyperparameters
 - different learning rate in backbone
 - other backbones
-- Pytorch Lightning
-- Better augmentation techniques. Cropping?
+- more complicated head-network
+- Pytorch Lightning or Keras implementations
+- Better augmentation techniques. Cropping? Filters?
 - Changes in dataset. Remove obvious crazy images.
+- Cross validation and other splits
+
+Dataset need to be structured like:
+ROOT > Classname > filename.jpg
 
 @author: Tom F. Hansen, Georg H. Erharter 
 """
@@ -29,35 +33,41 @@
 from torch.optim import lr_scheduler
 from torch.optim.optimizer import Optimizer
 from torchvision import datasets, models, transforms
+import torchvision
 
 import numpy as np
+import numpy.typing as npt
 from pathlib import Path
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import balanced_accuracy_score
+from sklearn.utils import class_weight
 import pickle
 from rich.traceback import install
 from rich.progress import track
 from typing import Tuple
+from utility import imshow
 
 # SETUPS
 ######################################################################################
 
 # presenting better error messages using rich
 install()
 
-ROOT = Path.cwd()
 DATA_DIR = Path(
     "/mnt/c/Users/TFH/NGI/TG Machine Learning - General/2022 ML workshop series/datasets/Rocks")
 TEST_SIZE = 0.3
 NUM_WORKERS = 12
+# remember to place both model and data on the same device
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+SAVE_PERFORMANCE = True
+SHOW_BATCH = True
 
 # Hyperparameters
 BATCH_SIZE = 64
 LR = 0.01
 MOMENTUM = 0.9
-STEP_SIZE = 3 #decay lr every xxx epoch
-GAMMA = 0.1 #decay factor for multiplication
+STEP_SIZE = 4 #decay lr every xxx epoch
+GAMMA = 0.3 #decay factor for multiplication
 NUM_EPOCHS = 10
 
 # stop randomness for model comparison
@@ -73,7 +83,7 @@
         transforms.RandomResizedCrop(224),
         transforms.RandomHorizontalFlip(),
         transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # values for transfer learning model
     ]),
     'test': transforms.Compose([
         transforms.Resize(256),
@@ -82,11 +92,12 @@
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ]),
 }
+
 # due to lazy_loading I don't allocate more memory of using the same dataset here
 train_dataset = datasets.ImageFolder(root=DATA_DIR, transform=data_transforms["train"])
 test_dataset = datasets.ImageFolder(root=DATA_DIR, transform=data_transforms["test"])
 
-# create splitting indices for samples - could also use SubsetRandomSampler
+# create splitting indices for samples
 num_classes = len(train_dataset.classes)
 indices = np.arange(len(train_dataset))
 labels = train_dataset.targets
@@ -96,15 +107,32 @@
 train_set = Subset(train_dataset,indices=train_ind)
 test_set = Subset(test_dataset,indices=test_ind)
 
+
+# Testing for comparison with the ants and bees dataset. You should get accuracies over 95% on that dataset.
+# Note that this dataset has just 245 images in training set and still the transfer learning model works well.
+# You can just uncomment this code and it should run, after you have updated with your path to the dataset
+# DATA_DIR_TRAIN = Path("/home/tfha/datasets/hymenoptera_data/train")
+# DATA_DIR_VAL = Path("/home/tfha/datasets/hymenoptera_data/val")
+
+# train_dataset = datasets.ImageFolder(root=DATA_DIR_TRAIN, transform=data_transforms["train"])
+# test_dataset = datasets.ImageFolder(root=DATA_DIR_VAL, transform=data_transforms["test"])
+# num_classes = len(train_dataset.classes)
+# train_set = train_dataset
+# test_set = test_dataset
+
+
 train_dataloader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True,num_workers=NUM_WORKERS)
 test_dataloader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False,num_workers=NUM_WORKERS)
 
-# sample = iter(train_dataloader).next()
-# sample[0] # batch of images
-# sample[1] # batch of labels
-# image1 = sample[0][0]
-# image1.shape
+#VISUALIZE A BATCH OF DATA
+######################################################################################
 
+if SHOW_BATCH:
+    # Get a batch of training data
+    inputs, classes = next(iter(train_dataloader))
+    # Make a grid from batch
+    out = torchvision.utils.make_grid(inputs)
+    imshow(out, title=[train_dataset.classes[x] for x in classes])
 
 # METHODS FOR TRAINING AND EVALUATION
 ######################################################################################
@@ -114,17 +142,17 @@ def train_epoch(
     optimizer: Optimizer, 
     loss_function: nn.CrossEntropyLoss, 
     dataloader: DataLoader, 
-    ):
+    )->Tuple[float, npt.NDArray, npt.NDArray]:
     """Train model for all samples in one epoch.
     Returning loss, labels, predictions"""
 
     epoch_loss = []
-    epoch_labels = np.array(())
-    epoch_preds = np.array(())
+    epoch_labels: npt.NDArray = np.array(())
+    epoch_preds: npt.NDArray = np.array(())
 
     # looping over all batches of samples
     for images, labels in track(dataloader,description="Training batches: "):
-        images = images.to(device)
+        images = images.to(device) #sending data to gpu or cpu
         labels = labels.to(device)
 
         logits = model(images) # forward pass
@@ -148,13 +176,13 @@ def test_epoch(
     model: nn.Module, 
     loss_function: nn.CrossEntropyLoss, 
     dataloader: DataLoader, 
-    )->Tuple[float, list, list]:
+    )->Tuple[float, npt.NDArray, npt.NDArray]:
     """Test model for all samples in one epoch.
     Returning loss, labels, predictions"""
 
     epoch_loss = []
-    epoch_labels = np.array(())
-    epoch_preds = np.array(())
+    epoch_labels: npt.NDArray = np.array(())
+    epoch_preds: npt.NDArray = np.array(())
 
     # looping over all batches of samples
     for images, labels in track(dataloader,description="Testing batches: "):
@@ -177,8 +205,9 @@ def test_epoch(
 # DEFINE NETWORK, LOSSFUNCTION, OPTIMIZER, LR-SCHEDULER
 ######################################################################################
 print(f"cuda is available: {torch.cuda.is_available()}. Device is {DEVICE}")
-model = models.resnet18(pretrained=True) #pretrained on 1000-class Imagenet database
-#turn of gradient update (learning) in backbone model
+model = models.resnet50(pretrained=True) #pretrained on 1000-class Imagenet database
+
+#turn off gradient update (learning) in backbone model
 for param in model.parameters():
     param.requires_grad = False
 
@@ -188,7 +217,12 @@ def test_epoch(
 
 model.to(DEVICE)
 
-loss_function = nn.CrossEntropyLoss()
+weights = class_weight.compute_class_weight(class_weight="balanced",
+                                                classes=np.unique(train_dataset.targets),
+                                                y=train_dataset.targets)
+weights = torch.tensor(weights).to(DEVICE)
+
+loss_function = nn.CrossEntropyLoss(weight=weights.float())
 optimizer = optim.SGD(model.fc.parameters(), lr=LR, momentum=MOMENTUM)
 LR_scheduler = lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)
 
@@ -199,46 +233,51 @@ def test_epoch(
 performance = []
 performance_path = Path("Reports/rock_classification_performance.pkl")
 tensorboard_path = Path("Reports/tensorboard_logs")
-writer = SummaryWriter(log_dir=tensorboard_path)
+writer = SummaryWriter(log_dir=tensorboard_path) # defines the Tensorboard writer
 
 for epoch in range(NUM_EPOCHS):
     # train
     model.train() # sets model in training mode. Turns on gradient update
     loss_train, train_labels, train_predictions = train_epoch(DEVICE,model, optimizer, loss_function,train_dataloader)
-    # acc_train = accuracy(train_predictions, train_labels, average="macro") # macro is balanced accuracy
     acc_train = balanced_accuracy_score(train_labels, train_predictions)
 
     # test
     model.eval() # Freeze model weights. No model update
     loss_test, labels, predictions = test_epoch(DEVICE, model, loss_function, test_dataloader)
     acc_test = balanced_accuracy_score(labels, predictions)
 
+    LR_scheduler.step()
+    
     # report metrics
-    print(f"Train-loss: {loss_train:.3f}. Train-acc: {acc_train:.2f}. \
-        Test-loss: {loss_test:.3f}. Test-acc: {acc_test:.3f}")
-
-    # add data to Tensorboard for live reporting
-    writer.add_scalars("Loss development",{
-        "Loss train": loss_train,
-        "Loss test": loss_test
-    }, global_step=epoch)
-    writer.add_scalars("Accuracy development",{
-        "Accuracy train": acc_train,
-        "Accuracy test": acc_test
-    }, global_step=epoch)
-
-    # save data to pickle file every epoch. Load data for later analysis
-    performance.append({ 
-    'epoch': epoch + 1,  #epoch counts from 0
-    'train_loss': loss_train,
-    'train_acc': acc_train,
-    'test_loss': loss_test,
-    'test_acc': acc_test,
-    'test_labels':labels,
-    'test_predictions':predictions
-    })
-    pickle.dump(performance, open(performance_path, 'wb'))
+    current_lr = LR_scheduler.get_last_lr()
+    print(f"Epoch: {epoch}. Train-loss: {loss_train:.3f}. Train-acc: {acc_train:.2f}. Test-loss: {loss_test:.3f}. Test-acc: {acc_test:.3f}. LR: {current_lr}")
+
+
+    if SAVE_PERFORMANCE:
+        # add data to Tensorboard for inspection of training development
+        writer.add_scalars("Loss development",{
+            "Loss train": loss_train,
+            "Loss test": loss_test
+        }, global_step=epoch)
+        writer.add_scalars("Accuracy development",{
+            "Accuracy train": acc_train,
+            "Accuracy test": acc_test
+        }, global_step=epoch)
+
+        # save data to pickle file every epoch. Load data for later analysis of results
+        performance.append({ 
+        'epoch': epoch + 1,  #epoch counts from 0
+        'train_loss': loss_train,
+        'train_acc': acc_train,
+        'test_loss': loss_test,
+        'test_acc': acc_test,
+        'test_labels':labels,
+        'test_predictions':predictions,
+        'class_names':train_dataset.classes
+        })
+        pickle.dump(performance, open(performance_path, 'wb'))
 
 writer.close()
 # save trained model for later predictions and analysis
-torch.save(model, Path("Reports/rock_model.pth"))
+if SAVE_PERFORMANCE:
+    torch.save(model, Path("Reports/rock_model.pth"))

src/utility.py

@@ -57,6 +57,18 @@ def plot_features(self, X: npt.NDArray)->None:
         plt.show()
 
 
+def imshow(inp, title=None):
+    """Imshow for Tensor."""
+    inp = inp.numpy().transpose((1, 2, 0))
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    inp = std * inp + mean
+    inp = np.clip(inp, 0, 1)
+    plt.imshow(inp)
+    if title is not None:
+        plt.title(title)
+    plt.pause(0.001)  # pause a bit so that plots are updated
+    
 
 def examplify():
     """Examplifies functionality"""