I’m working on an image classification model for age recognition using the Fair Face dataset, and after loading in the dataset and checking if everything works correctly, and then loading it all into data loaders, the code doesn’t iterate through the dataset for training. There are also no error messages.

I’ve tried entering print statements into the training loop to see where the training was having problems, but apparently once it starts enumerating through the dataset, it freezes at the beginning of the first batch.

Does anyone know why this is happening, and how it can be fixed?

Here’s the custom dataset class:

class FairFaceDataset(Dataset):
    """Face Landmarks dataset."""

    def __init__(self, root_dir, csv_file, transform=None):
        """
        Arguments:
            csv_file (string): Path to the csv file with the labels.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.csv = self.data = pd.read_csv(csv_file).iloc[:, 1]
        self.dir = root_dir
        self.image_files = os.listdir(root_dir)
        self.transform = transform

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

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        img_name = os.path.join(self.dir,
                                self.image_files[idx])
        image = io.imread(img_name)
        label = self.csv[idx]

        if self.transform:
            image = self.transform(image)

        return (image, label)

The loading of the data:

transform_img = transforms.Compose([transforms.ToTensor(), 
                                    #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                   ])

# set batch_size
batch_size = 4

# set number of workers
num_workers = 2

trainset = FairFaceDataset(train_dir, train_csv, transform_img)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=num_workers)

testset = FairFaceDataset(test_dir, test_csv, transform_img)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=num_workers)

race_scores_fair = ["White", "Black", "Latin Hispanic", "East", "Southeast_Asian", "Indian", "Middle Eastern"]
race_scores_fair_4 = ["White", "Black", "Asian", "Indian"]
gender_scores_fair = ["Male", "Female"]
age_scores_fair = ["0-2", "3-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70+"]

The network:

class Net(nn.Module):
    ''' Models a simple Convolutional Neural Network'''

    def __init__(self):
      ''' initialize the network '''
      super(Net, self).__init__()
      self.dropout = nn.Dropout(0.2)
        # 1 input image channel, 6 output channels,
        # 5x5 square convolution kernel
      self.conv1 = nn.Conv2d(3, 12, 3)
        # Max pooling over a (2, 2) window
      self.pool = nn.MaxPool2d(2, 2)
      self.conv2 = nn.Conv2d(12, 24, 3)

      self.fc1 = nn.Linear(25 * 5 * 5, 300)
      self.fc2 = nn.Linear(300, 150)
      self.fc3 = nn.Linear(150, 9)

    def forward(self, x):
      ''' the forward propagation algorithm '''
      x = self.pool(F.relu(self.conv1(x)))
      x = self.dropout(x)
      x = self.pool(F.relu(self.conv2(x)))
      x = self.dropout(x)
      x = x.flatten(1)
      x = F.relu(self.fc1(x))
      x = F.relu(self.fc2(x))
      x = F.relu(self.fc3(x))
      return x

net = Net()
print(net)

And the training loop:

print(torch.cuda.is_available())

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()

for epoch in range(25):  # loop over the dataset multiple times
    print("I'm an epoch!")

    running_loss = 0.0
    for i, data in enumerate(trainloader):
        print("1")
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data
        print("2")

        # zero the parameter gradients
        optimizer.zero_grad()
        print("3")

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        print("4")

        # print statistics
        running_loss += loss.item()
        if i % 500 == 499:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch, i + 1, running_loss / 500))
            running_loss = 0.0

    before_lr = optimizer.param_groups[0]["lr"]
    #scheduler.step()
    after_lr = optimizer.param_groups[0]["lr"]

    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('   Accuracy -> %.2f' % (
        100 * correct / total), end="")
    print("%")


# whatever you are timing goes here
end.record()

# Waits for everything to finish running
torch.cuda.synchronize()

print('Finished Training')
print(start.elapsed_time(end))  # milliseconds

I already went through all the code, and looked at the PyTorch website, and other tutorials for how they iterated through custom datasets, but I couldn’t find anything when it comes to specifically to training on custom datasets. I think this is because after it’s loaded into the data loader everything should work as normal.