I am trying to train a convolutional neural network on an image classification task. For some reason my dataloader class is picking up the parent folder and I suspect this is causing issues when attempting to train the model as it is giving me an error “RuntimeError: weight tensor should be defined either for all 64 classes or no classes but got weight tensor of shape: [99]”
Here’s my code:
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader, Dataset, random_split
import torchvision
from torchvision import datasets, models, transforms
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image
data_path="/kaggle/input/facial-age/face_age"
class CustomImageFolder(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.images = []
self.labels = []
# Iterate over all subdirectories
for label in sorted(os.listdir(root_dir)):
if label == 'face_age':
continue
label_path = os.path.join(root_dir, label)
if os.path.isdir(label_path):
for img_file in os.listdir(label_path):
img_path = os.path.join(label_path, img_file)
if os.path.isfile(img_path):
self.images.append(img_path)
self.labels.append(label)
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
img_path = self.images[idx]
image = Image.open(img_path).convert('RGB')
label = self.labels[idx]
if self.transform:
image = self.transform(image)
return image, int(label)
# Define the transform
transform = transforms.Compose([transforms.ToTensor()])
# Create the custom dataset
dataset = CustomImageFolder(root_dir=data_path, transform=transform)
# Check the first few items
for i in range(5):
image, label = dataset[i]
print(f'Label: {label}, Image Shape: {image.shape}')
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
dataset = datasets.ImageFolder(root=data_path)
print(dataset.class_to_idx)
print(len(dataset.class_to_idx))
The last cell is outputting the parent folder along with all the folder inside it. This makes no sense.
here’s the dataset used: https://www.kaggle.com/datasets/frabbisw/facial-age