I've implemented AlexNet() and train_model(). I used a dataset of three classes named Dog, Food and Vehicles. After pre_processing and using test_train split() , the Data I'm feeding to the train_model is throwing an error . Check what went wrong . I think the data I'm passing and the data that train_model() is expecting and the data ALexnet are all different altogether causing serious error. FIX THIS ONCE AND FOR ALL. If possible try not to use cv2. You can use keras/tensorflow or pytorch , sklearn. Do not use any in built fit functions.        The code for preprocessing the images should be as light as possible and also Provide three visualizations. the implemented visulisatiions are throwing an error and images resolutions displayed is not so good.                                                                                         import torch,cv2
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
import numpy as np
import pandas as pd
from concurrent.futures import ThreadPoolExecutor

def process_image(file_path, size=(64, 64)):
    img = cv2.imread(file_path)
    img = cv2.resize(img, size)
    return np.array(img, dtype=np.float32) / 255.0

def load_and_preprocess(folder_names):
    folder_paths = [os.path.join(folder_name, file) for folder_name in folder_names for file in os.listdir(folder_name)]
    labels = [folder_name for folder_name in folder_names for _ in os.listdir(folder_name)]

    with ThreadPoolExecutor() as executor:
        processed_images = list(executor.map(process_image, folder_paths))

    return np.array(processed_images), pd.get_dummies(labels), labels


X, y, label_list = load_and_preprocess(folder_names)



# Create train and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify = y)



class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()

        self.features = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )

        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, 1000),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), 256 * 6 * 6)
        x = self.classifier(x)
        return x

alexnet = AlexNet()



and my train_model implementation:


from sklearn.metrics import accuracy_score

def train_model(model, X_train, y_train, X_test, y_test, epochs = None, batch_size = None, learning_rate = None, optimization_technique = None,  patience=None, scheduler_patience=None, **kwargs):
    # Define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)

    if optimization_technique == 'learning_rate_scheduler' and scheduler_patience:
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=scheduler_patience, verbose=True)

    # Training segment
    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []

    best_loss = float('inf')
    stopping_counter = 0

    for epoch in range(epochs):
        epoch_train_losses = []
        epoch_y_true = []
        epoch_y_pred = []

        for i in range(0, len(X_train), batch_size):
            X_batch = torch.tensor(X_train[i:i + batch_size], dtype=torch.float).permute(0, 3, 1, 2)
            y_batch = torch.tensor(y_train[i:i + batch_size].values, dtype=torch.long).view(-1)

            optimizer.zero_grad()
            y_pred = model(X_batch)
            loss = criterion(y_pred, y_batch)
            loss.backward()
            optimizer.step()

            epoch_train_losses.append(loss.item())
            epoch_y_true.extend(y_batch.numpy().flatten().tolist())
            epoch_y_pred.extend(y_pred.argmax(dim=1).numpy().flatten().tolist())

        train_losses.append(sum(epoch_train_losses) / len(epoch_train_losses))
        train_accuracies.append(accuracy_score(epoch_y_true, epoch_y_pred))

        # Testing segment
        with torch.no_grad():
            X_test_tensor = torch.tensor(X_test, dtype=torch.float32).permute(0, 3, 1, 2)
            y_test_tensor = torch.tensor(y_test.values, dtype=torch.long).view(-1)

            test_pred = model(X_test_tensor)
            _, preds_indices = torch.max(test_pred, 1)
            test_loss = criterion(test_pred, y_test_tensor)
            test_accuracy = accuracy_score(y_test_tensor, preds_indices)

        test_losses.append(test_loss.item())
        test_accuracies.append(test_accuracy)

        # Early stopping
        if optimization_technique == 'early_stopping' and patience:
            if test_loss < best_loss:
                best_loss = test_loss
                stopping_counter = 0
            else:
                stopping_counter += 1
            if stopping_counter > patience:
                print(f"Early stopping at epoch {epoch+1}/{epochs}")
                break

        # Learning rate scheduler
        if optimization_technique == 'learning_rate_scheduler' and scheduler_patience and scheduler:
            scheduler.step(test_loss)

        if epoch % 100 == 0:
            print(f"Epoch: {epoch+1}/{epochs}, Training Loss: {train_losses[-1]}, Test Loss: {test_loss.item()}, Training Accuracy: {train_accuracies[-1]}, Test Accuracy: {test_accuracy}")

    return train_losses, train_accuracies, test_losses, test_accuracies