Transfer Learning with CNNs

You will fine-tune the pre-trained ResNet-18 network which is available at Pytorch. This network is trained on ImageNet dataset so you are not initializing the weights randomly, instead, you are using the pre-trained weights from this network. Freeze all the layers before training the network, except the FC layer. Consequently, the gradients will not be calculated for the layers except the ones that you are going to update (FC layer) over the pre-trained weights. However, since the number of classes is different for our dataset, you should modify the last layer of the network, which the probabilities will be calculated on. Therefore, the weights will be randomly initialized for this layer.
1. What is fine-tuning? Why should we do this? Why do we freeze the rest and train only FC layers? Give your
explanation in detail with relevant code snippet. 
2. Explore training with two different cases; train only FC layer and freeze rest, train last two convolutional layers and FC layer and freeze rest. Tune your parameters accordingly and give accuracy on validation set and test
set. Compare and analyze your results. Give relevent code snippet.
3. Plot confusion matrix for your best model and analyze results.

Here is the code provided to you including data loading and plotting results:
"
def plot(train_losses, valid_losses, valid_accuracies, title):
    epochs = range(1, len(train_losses) + 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_losses, label="Training Loss")
    plt.plot(epochs, valid_losses, label="Validation Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochs, valid_accuracies, label="Validation Accuracy")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()

    plt.suptitle(title)  # Add the title here

    timestamp = time.strftime("%Y%m%d%H%M%S")
    filename = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), f"image_out/plot_{timestamp}.png"
    )

    # Save the plot with the unique filename
    plt.savefig(filename)

    plt.show()


def plotTest(test_loss, test_accuracy):
    epochs = range(1, len(test_loss) + 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, test_loss, label="Test Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochs, test_accuracy, label="Test Accuracy")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()

    plt.show()


def plotDropOut(valid_losses, valid_accuracies, test_loss, test_accuracy):
    epochs = range(1, len(train_losses) + 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, valid_losses, label="Validation Loss")
    plt.plot(epochs, test_loss, label="Test Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(epochs, valid_accuracies, label="Validation Accuracy")
    plt.plot(epochs, test_accuracy, label="Test Accuracy")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()

    plt.show()


def plot_confusion_matrix(y_true, y_pred, labels):
    cm = confusion_matrix(y_true, y_pred, labels=labels)
    sns.heatmap(
        cm, annot=True, cmap="Blues", fmt="d", xticklabels=labels, yticklabels=labels
    )
    plt.xlabel("Predicted")
    plt.ylabel("True")
    plt.show()


# -------
# Dataset
# -------


class BacteriaDataset(Dataset):
    def __init__(self, data, transform=None):
        self.data = data
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        image, label = self.data[index]
        if self.transform:
            image = self.transform(image)
        return image, label


def load_data(data_dir, transform, min_images_per_class=50, min_test_per_class=20):
    categories = os.listdir(data_dir)
    data = []
    class_counts = {}  # Track the number of images per class
    for label, category in enumerate(categories):
        class_counts[label] = 0  # Initialize the count for each class
        category_dir = os.path.join(data_dir, category)
        for img_name in os.listdir(category_dir):
            img_path = os.path.join(category_dir, img_name)
            img = Image.open(img_path).convert("RGB")
            img = img.resize((128, 128))
            data.append((img, label))
            class_counts[label] += 1  # Increment the count for each image added

    # Shuffle the data once for random splitting
    random.shuffle(data)

    # Ensure that each split has at least the required number of images per class

    train_data, valid_data, test_data = [], [], []
    class_counts_train, class_counts_valid, class_counts_test = {}, {}, {}

    for label in class_counts:
        class_counts_train[label] = 0
        class_counts_valid[label] = 0
        class_counts_test[label] = 0

    for img, label in data:
        if class_counts_train[label] < min_images_per_class:
            train_data.append((img, label))
            class_counts_train[label] += 1
        elif class_counts_valid[label] < min_test_per_class:
            valid_data.append((img, label))
            class_counts_valid[label] += 1
        elif class_counts_test[label] < min_test_per_class:
            test_data.append((img, label))
            class_counts_test[label] += 1
        else:
            test_data.append((img, label))

    return train_data, valid_data, test_data

data_dir = os.path.join(os.getcwd(), "Micro_Organism") # Setting path to the dataset folder
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

transform = transforms.Compose([
    transforms.Resize((64, 64)),
    # transforms.RandomHorizontalFlip(),
    # transforms.RandomRotation(20),
    transforms.ToTensor(),
])

train_data, valid_data, test_data = load_data(data_dir, transform)

train_set = BacteriaDataset(train_data, transform)
valid_set = BacteriaDataset(valid_data, transform)
test_set = BacteriaDataset(test_data, transform)

train_loader = DataLoader(train_set, batch_size=bs, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=bs, shuffle=False)
test_loader = DataLoader(test_set, batch_size=best_model_info["bs"], shuffle=False)
"