Training DataLoader, Loop does not iterate

Hi everyone, can someone help me out?

I am doing a project that consists in create a CNN which identifies if the image has cancer or not.

The problem is when I try to train the model, I can’t iterate the the loader. The error that it does not say anything.

This is the code:

# Training function
def train_mri(model: MRIClassifier, train_loader: DataLoader, criteration: torch.nn.BCEWithLogitsLoss, optimizer: optim.Adam, num_epochs: int):
    # Set the model to training mode
    model.train()

    for epoch in range(num_epochs):
        runnin_loss = 0.0
        correct = 0
        total = 0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.float().to(device) # Move to GPU
            optimizer.zero_grad() # Zero the Gradients
            outputs = model(images) # Forward pass
            loss = criteration(outputs, labels.unsqueeze(1)) # Compute Loss
            loss.backward() # Backward Pass
            optimizer.step() # Optimize Weights

            runnin_loss += loss.item() # Accumulate Loss

            # Calculate Accuracy
            predicted = (torch.sigmoid(outputs) > 0.5).float() # Get predicted classes
            total += labels.size(0)
            correct += (predicted == labels.unsqueeze(1)).sum().item() # Count correct predictions

        # Calculate average loss and accuracy
        avg_loss = runnin_loss / len(train_loader)
        accuracy = 100 * correct / total
        print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}, Accuracy {accuracy:.4f}%")

model.to(device)

# Define the loss function and optimizer
criteration = torch.nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Train the model
train_mri(model, train_loader, criteration, optimizer, num_epochs=20)

This is the error:

KeyError                                  Traceback (most recent call last)
File c:\Users\asus\miniconda3\envs\Deep-Learning\lib\site-packages\pandas\core\indexes\base.py:3805, in Index.get_loc(self, key)
   3804 try:
-> 3805     return self._engine.get_loc(casted_key)
   3806 except KeyError as err:

File index.pyx:167, in pandas._libs.index.IndexEngine.get_loc()

File index.pyx:196, in pandas._libs.index.IndexEngine.get_loc()

File pandas\\_libs\\hashtable_class_helper.pxi:7081, in pandas._libs.hashtable.PyObjectHashTable.get_item()

File pandas\\_libs\\hashtable_class_helper.pxi:7089, in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 9

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Cell In[72], line 8
      5 optimizer = optim.Adam(model.parameters(), lr=0.001)
      7 # Train the model
----> 8 train_mri(model, train_loader, criteration, optimizer, num_epochs=20)
...
   3815     #  InvalidIndexError. Otherwise we fall through and re-raise
   3816     #  the TypeError.
   3817     self._check_indexing_error(key)

KeyError: 9

The error shows a KeyError in pandas so check if you are indexing the DataFrame in an invalid way.

I did some transformations in the data, like it is shown in the following code:

def transform_data(df: pd.DataFrame, transform: torchvision.transforms.Compose):
    """
    Applies a transformation to a dataset of images and returns the transformed dataset.

    Args:
        dataset (pd.DataFrame): DataFrame containing image file paths in the first column and labels in the second column.
        transform (torchvision.transforms.Compose): Transformations to be applied to the images.

    Returns:
        pd.DataFrame: A DataFrame containing the transformed images as tensors and their corresponding labels.
    """

    images = [] # List to store transformed image tensors
    labels = [] # List to store labels

    # Iterate through each image path and label in the dataset
    for img_path, label in zip(df.iloc[:, 0], df.iloc[:, 1]):
        img = Image.open(img_path) # Open the image file

        transformed_img = transform(img) # Apply the transformation

        images.append(transformed_img) # Store the transformed image tensor
        labels.append(label) # Store the corresponding label

    # Create a new DataFrame with transformed images and labels
    transformed_df = pd.DataFrame({
        "Image": images,
        "Label": labels
    })

    return transformed_df.reset_index(drop=True)

And then I made another transformations after this one, only for the training set:

def train_augmentation(train_df: pd.DataFrame, transform: torchvision.transforms.Compose, num_augmentations=1) -> pd.DataFrame: 
    """
    Augments a dataset by applying transformations to each image.

    Args:
        df (pd.DataFrame): Training DataFrame containing images and their labels.
        transform (torchvision.transforms.Compose): Transformations to apply to the images.
        num_augmentations (int): Number of augmented versions to generate for each image.

    Returns:
        pd.DataFrame: A new DataFrame containing both original and augmented images with their corresponding labels.
    """

    images = []
    labels = []

    # Iterate through the dataset to process each image and label
    for image, label in zip(train_df.iloc[:, 0], train_df.iloc[:, 1]):
        # Add the original image and label
        images.append(image)
        labels.append(label)

        # Apply transformations to create augmented images
        for _ in range(num_augmentations):
            transform_img = transform(image)
            images.append(transform_img)
            labels.append(label)
    
    # Create a new DataFrame with augmented images and labels
    augmented_dataframe = pd.DataFrame({
        "Images": images,
        "Labels": labels
    })

    return augmented_dataframe.reset_index(drop=True)

After every transformation I passed the data to DataLoaders:

train_loader = DataLoader(dataset=train_augmented, batch_size=16, shuffle=True)
val_loader = DataLoader(dataset=val_df, batch_size=16, shuffle=True)
test_loader = DataLoader(dataset=test_df, batch_size=16, shuffle=True)

train_loader, val_loader, test_loader

I used reset_index(drop=True) function, to reset the index and remove the previous one.
And the error still appears.