Skip to main content
In this lesson, we will guide you through building and preprocessing a custom dataset to train a machine learning model. Whether you are working with images, text, audio, or any other data modality, this step-by-step tutorial covers data cleaning, annotation creation, dataset splitting, versioning, applying transformations, and ultimately preparing PyTorch datasets along with DataLoaders.

1. Loading and Displaying the Dataset

First, we load the dataset and visualize the images to verify that they meet the training requirements.
Viewing your dataset before training helps identify any images that do not belong to your target classes.
Next, we display each image along with its file name. This step assists in detecting any out-of-scope images.
In our dataset, the expected images are of cats and dogs. However, the initial exploration might reveal images of a frog or horse, which are irrelevant for a cat-versus-dog classification task.

2. Cleaning the Dataset and Creating Annotations

After examining the images, it’s crucial to remove any that do not match the target classes. In this example, we remove the horse and frog images from the dataset and then generate an annotations CSV file.
Now, create the annotations CSV file where each record maps the image file path to its class label (extracted from the directory structure).
The resulting CSV (image_data.csv) should look similar to: file_path,label
images/cat/cat-4.jpg,cat
images/cat/cat-5.jpg,cat
images/cat/cat-2.jpg,cat
images/cat/cat-3.jpg,cat
images/dog/dog-4.jpg,dog
images/dog/dog-1.jpg,dog
images/dog/dog-3.jpg,dog
images/dog/dog-5.jpg,dog

3. Creating an Initial PyTorch Dataset

Next, we create an initial PyTorch dataset class that reads our annotations CSV file and returns the image path along with its label. This forms the basis for our training pipeline.

4. Splitting the Dataset

It is important to split the dataset into training, validation, and testing subsets. Here, we randomly partition the data into 70% for training, 15% for validation, and 15% for testing.
The indices printed represent the positions of images in the original dataset. You can further verify the allocation by printing rows from dataset.img_labels:

5. Data Versioning and Annotation for Subsets

Versioning your data annotations is a best practice for reproducibility. By saving separate CSV files for each subset (training, validation, and testing), you can easily track and reproduce your training experiments. For example, to generate annotations for the training set:
Similarly, create annotation files for the validation and testing sets:
This separation helps maintain clear records of which images are used during each phase of training.

6. Defining Data Transformations

Data transformations and augmentations are key to preparing your images for model training. Typically, training data benefits from a variety of augmentations, while validation data should remain consistent.

Training Transformations

In this example, we use TorchVision’s v2 transforms to resize images, perform random cropping and horizontal flipping, convert to tensors, and apply normalization.

Validation Transformations

For validation, we avoid random augmentations to ensure consistent inputs.

7. Constructing a Custom Image Dataset and DataLoaders

We now build a custom dataset class that incorporates our annotations, image directory, and transformation pipelines. Additionally, we use a label encoding strategy to convert categorical labels into numerical format.
Create the validation dataset similarly:

Creating DataLoaders

DataLoaders help batch and shuffle data during training and evaluation.
Due to the use of random cropping in the training transformations, the spatial dimensions of training images (e.g., 75×75) might differ from the fixed dimensions of the validation images (128×128).

Conclusion

Congratulations! You have now learned how to:
  • Load and visualize image data.
  • Clean the dataset and create annotation CSVs.
  • Build an initial PyTorch dataset and split it into training, validation, and testing subsets.
  • Implement data versioning for reproducibility.
  • Define and apply transformation pipelines for data augmentation.
  • Develop custom PyTorch datasets and DataLoaders.
With these foundational steps, you are well-equipped to proceed with model training using your custom data. For further guidance, consider reviewing resources like PyTorch Documentation and TorchVision Transforms. Happy coding!

Watch Video

Practice Lab