Dataset From Files

Real-world projects often require working with raw image files rather than prepackaged datasets. In this tutorial, we will download a cats-versus-dogs image archive, load images from class-specific folders, apply deterministic preprocessing, and divide the data into separate training, validation, and test sets.

Show code cell source

Hide code cell source

import hashlib
import os
import pathlib
import shutil
import tarfile
import zipfile

import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import torchvision.transforms.v2 as v2
from PIL import Image
from sklearn.model_selection import train_test_split
from torchvision.datasets import ImageFolder
from tqdm import tqdm

Data preparation

We use Microsoft’s archived copy of the cats-vs-dogs dataset that was made available by Kaggle as part of a computer vision competition in 2013. This dataset contains roughly 25’000 JPEG images of dogs and cats, for a total size of 787 MB. The fixed URL identifies the dataset snapshot used here; in a production pipeline, also record a trusted SHA-256 checksum or dataset version in experiment metadata.

Show code cell source

Hide code cell source

def _safe_members(archive, destination):
    """Reject archive members that would escape the destination directory."""
    root = pathlib.Path(destination).resolve()
    for member in archive.getmembers() if isinstance(archive, tarfile.TarFile) else archive.infolist():
        name = member.name if isinstance(archive, tarfile.TarFile) else member.filename
        target = (root / name).resolve()
        if root not in target.parents and target != root:
            raise ValueError(f"Unsafe archive member: {name}")
        if isinstance(archive, tarfile.TarFile) and (member.issym() or member.islnk()):
            raise ValueError(f"Archive links are not allowed: {name}")
        yield member


def download_file(url, save_dir, expected_sha256=None):
    save_dir = pathlib.Path(save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    filepath = save_dir / pathlib.Path(url).name

    if not filepath.exists():
        torch.hub.download_url_to_file(url, filepath)
    else:
        print(f"Using cached file: {filepath}")

    if expected_sha256 is not None:
        with filepath.open("rb") as stream:
            digest = hashlib.file_digest(stream, "sha256").hexdigest()
        if digest != expected_sha256:
            raise ValueError(f"Checksum mismatch for {filepath}")

    print(f"Extracting: {filepath}")
    if filepath.suffix == ".zip":
        with zipfile.ZipFile(filepath) as archive:
            archive.extractall(save_dir, members=_safe_members(archive, save_dir))
    elif tarfile.is_tarfile(filepath):
        with tarfile.open(filepath, "r:*") as archive:
            archive.extractall(save_dir, members=_safe_members(archive, save_dir))

url = "https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_5340.zip"
folder = ".data/cats_vs_dogs"

download_file(url, folder)

data_dir = pathlib.Path(folder) / "PetImages"

100%|██████████| 787M/787M [00:11<00:00, 74.6MB/s] 

Extracting: .data\cats_vs_dogs\kagglecatsanddogs_5340.zip

Handling invalid images

Large, real-world datasets commonly contain truncated or incorrectly encoded files. We validate each image with PIL before constructing the dataset. We move corrupted images to a quarantine directory and record the reason. This keeps the cleanup auditable and makes it possible to inspect or restore a file later.

def is_valid_image(path):
    import warnings
    try:
        with warnings.catch_warnings():
            warnings.simplefilter("error")

            with Image.open(path) as image:
                image.verify()

            with Image.open(path) as image:
                image.load()
                
        return True, None
    
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"

We now validate the JPEG files and quarantine any that PIL cannot decode reliably.

Show code cell source

Hide code cell source

image_paths = list(data_dir.rglob("*.jpg"))
quarantine_dir = pathlib.Path(folder) / ".invalid_images"
invalid_images = []

with tqdm(image_paths, desc="Validating images", unit=" image", dynamic_ncols=True) as progress:
    for path in progress:
        valid, reason = is_valid_image(path)
        if not valid:
            relative_path = path.relative_to(data_dir)
            destination = quarantine_dir / relative_path
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(path, destination)
            invalid_images.append((str(relative_path), reason))
            progress.set_postfix(quarantined=len(invalid_images))

print(f"Quarantined {len(invalid_images)} images in {quarantine_dir}.")

Validating images: 100%|██████████| 25000/25000 [02:21<00:00, 176.43 image/s, quarantined=3] 

Quarantined 3 images in .data\cats_vs_dogs\.invalid_images.

Dataset ImageFolder

Currently, the images sit on a drive as JPEG files. They must be loaded into memory and preprocessed before being passed to a neural network. This involves the following steps.

• Read a picture file.

• Decode the JPEG content to RGB grids of pixels.

• Convert it into a floating-point tensor properly preprocessed.

TorchVision provides a utility class called ImageFolder that does all this for us. It is designed for loading an image dataset organized in a specific directory structure. Within the root directory, each subfolder represents a class, and the images contained in that folder are treated as belonging to that class. For example, in the cats-vs-dogs dataset, the root directory should have two subfolders. The Cat subfolder contains images that belong to the “cat” class, and the Dog subfolder contains images that belong to the “dog” class. The structure looks like the following.

Dataset/
│
├── Cat/
│   ├── cat001.jpg
│   ├── cat002.jpg
│   └── ...
│
└── Dog/
    ├── dog001.jpg
    ├── dog002.jpg
    └── ...

We will use the ImageFolder class to load the cats-vs-dogs dataset. But first, we need to define the preprocessing steps that we will apply to the images.

Preprocessing

The ImageFolder dataset accepts a transform for images and an optional target_transform for labels. Our deterministic preprocessing converts each PIL image to a tensor, resizes its shorter edge, takes a centered square crop, and rescales integer pixel values from [0, 255] to floating-point values in [0, 1].

preprocess = v2.Compose([
    v2.ToImage(),
    v2.Resize(128),
    v2.CenterCrop(128),
    v2.ToDtype(torch.float32, scale=True),
])

We have already seen the transformations ToImage() and ToDtype() in a previous tutorial. Let’s take a moment to understand the transformations Resize() and CenterCrop() added to the pipeline.

Note

For a model trained from scratch, pixel scaling is a reasonable baseline. Pretrained models normally also expect channel-wise normalization with specific mean and standard deviation.

Resize

The Resize transform modifies the size of the input image. The first argument size is mandatory and can be an integer, a tuple, or None. The behavior of the transform depends on the value passed to this argument.

• If size is a tuple (height, width), the output size will be matched to this tuple.

• If size is an integer, the smaller edge of the image will be matched to this number, and the larger edge will be scaled to maintain the aspect ratio. For example, if height > width, the image will be rescaled to (size * height / width, size).

• If size is None, the optional argument max_size must be set to an integer. Then the larger edge of the image will be matched to this number, and the smaller edge will be scaled to maintain the aspect ratio. For example, if height > width, the image will be rescaled to (max_size, max_size * width / height).

Let’s see an example of how the Resize transform works.

Show code cell source

Hide code cell source

original = Image.open(data_dir / "Cat" / "2.jpg")
resized1 = v2.Resize(128)(original)
resized2 = v2.Resize((128,128))(original)

plt.figure(figsize=(8, 4))
plt.subplot(1, 3, 1)
plt.imshow(original)
plt.title("Original")
plt.subplot(1, 3, 2)
plt.imshow(resized1)
plt.title("Resize(128)")
plt.subplot(1, 3, 3)
plt.imshow(resized2)
plt.title("Resize((128,128))")
plt.show()

Center Crop

The CenterCrop transform crops the input image at the center.

• If the size argument is an integer, the output size will be (size, size).

• If the size argument is a tuple (height, width), the output size will be matched to this.

Let’s see an example of how the CenterCrop transform works. Note that the combination of Resize and CenterCrop works well to resize the image to a square shape without changing the aspect ratio.

Show code cell source

Hide code cell source

original = Image.open(data_dir / "Cat" / "2.jpg")
crop1 = v2.CenterCrop(128)(original)
crop2 = v2.Compose([v2.Resize(128), v2.CenterCrop(128)])(original)

plt.figure(figsize=(8, 4))
plt.subplot(1, 3, 1)
plt.imshow(original)
plt.title("Original")
plt.subplot(1, 3, 2)
plt.imshow(crop1)
plt.title("Center Crop")
plt.subplot(1, 3, 3)
plt.imshow(crop2)
plt.title("Resize + Center Crop")
plt.show()

And more…

There are many other transformations available in TorchVision. This example provides a visual illustration of several transformations. For the complete list of transformations, please refer to the documentation.

Splitting the dataset

Because the cats-vs-dogs archive has no predefined split, we create three non-overlapping subsets.

• Training set for updating model parameters;

• Validation set for monitoring generalization and selecting a checkpoint;

• Test set for providing a final estimate after all choices are fixed.

We stratify both split operations so the cat/dog proportions remain similar and use a fixed random seed for reproducibility. Here we assume that the images are independent and identically distributed (i.i.d.). In practice, this assumption may not hold. For example, if multiple images of the same animal appear in the dataset, a random image-level split can leak information across subsets.

Note

There are several ways to split a dataset in PyTorch.

• Use the function random_split to randomly split the dataset into non-overlapping subsets.

• Use the class Subset to create a new dataset from a subset of elements of another dataset.

• Manually separate the image files into one directory per subset, and use the ImageFolder class to load each subset separately.

# Load the full dataset
dataset = ImageFolder(data_dir, transform=preprocess)

all_idx = list(range(len(dataset)))

# Reserve 70% for training
train_idx, holdout_idx = train_test_split(
    all_idx,
    stratify=dataset.targets,
    test_size=0.30,
    random_state=42,
)

holdout_targets = [dataset.targets[i] for i in holdout_idx]

# Reserve 15% for validation and 15% for testing
valid_idx, test_idx = train_test_split(
    holdout_idx,
    stratify=holdout_targets,
    test_size=0.50,
    random_state=42,
)

train_ds = torch.utils.data.Subset(dataset, train_idx)
valid_ds = torch.utils.data.Subset(dataset, valid_idx)
test_ds  = torch.utils.data.Subset(dataset, test_idx)

Let’s print some information about the datasets to verify that everything is working as expected.

Show code cell source

Hide code cell source

def class_percentages(indices):
    labels = torch.tensor([dataset.targets[i] for i in indices])
    return torch.bincount(labels, minlength=len(dataset.classes)) / len(indices) * 100


for name, indices in [
    ("All", all_idx),
    ("Train", train_idx),
    ("Valid", valid_idx),
    ("Test", test_idx),
]:
    percentages = class_percentages(indices)
    distribution = ", ".join(
        f"{label}: {value:.1f}%" for label, value in zip(dataset.classes, percentages)
    )
    print(f"{name:5s}: {len(indices):5d} | {distribution}")

All  : 24997 | Cat: 50.0%, Dog: 50.0%
Train: 17497 | Cat: 50.0%, Dog: 50.0%
Valid:  3750 | Cat: 50.0%, Dog: 50.0%
Test :  3750 | Cat: 50.0%, Dog: 50.0%

Visualization

As a final data check, inspect examples from the training and validation sets along with their numeric labels. The printed class_to_idx mapping above tells us what those labels mean. We avoid browsing the test set so it remains a clean final evaluation sample.

Show code cell source

Hide code cell source

def show_images(dataset, n=6, prefix=""):
    images, labels = zip(*[dataset[i] for i in range(n)])
    side_by_side = torch.cat(images, dim=2).permute(1, 2, 0)

    plt.figure(figsize=(10, 4))
    plt.imshow(side_by_side)
    plt.title(prefix + " images with labels: " + str(labels))
    plt.axis("off")
    plt.show()


show_images(train_ds, prefix="Train")
show_images(valid_ds, prefix="Validation")

Baseline model

We will train a simple convolutional network on the cats-vs-dogs dataset. The model consists of several convolutional layers mixed with max-pooling layers, followed by a fully-connected layer and a single-output layer for binary classification. The sigmoid activation is not included in the model because we will use nn.BCEWithLogitsLoss, which combines the sigmoid activation and the binary cross-entropy loss.

class BaselineModel(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        ksize = 3
        self.conv1 = torch.nn.Conv2d(3, 32, ksize)
        self.conv2 = torch.nn.Conv2d(32, 64, ksize)
        self.conv3 = torch.nn.Conv2d(64, 128, ksize)
        self.conv4 = torch.nn.Conv2d(128, 128, ksize)
        flat_dim = self.__calc_dim(128) # Assume 128x128 images
        self.fc1 = torch.nn.Linear(flat_dim, 256)
        self.fc2 = torch.nn.Linear(256, 1)

    def forward(self, x):
        for conv in [self.conv1, self.conv2, self.conv3, self.conv4]:
            x = F.relu(conv(x))
            x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
    
    def __calc_dim(self, input_dim: int):
        """Returns the tensor size after the flatten layer."""
        n = input_dim
        for conv in [self.conv1, self.conv2, self.conv3, self.conv4]:
            n = n - conv.kernel_size[0] + 1 # CONV output size
            n = n // 2                      # POOL output size
        dim = n * n * self.conv4.out_channels
        return dim

Let’s take a look at how the dimensions of the feature maps change with every successive layer. Here, since we start from inputs of size 128x128 (a somewhat arbitrary choice), we end up with feature maps of size 6x6 right before the flatten layer. Note that the depth of the feature maps is progressively increasing in the network (from 32 to 128), while the size of the feature maps is decreasing (from 126x126 to 6x6). This is a pattern that you will see in almost all convnets.

Layer	Output Shape
Input	(…, 3, 128, 128)
Conv1	(…, 32, 126, 126)
Pool1	(…, 32, 63, 63)
Conv2	(…, 64, 61, 61)
Pool2	(…, 64, 30, 30)
Conv3	(…, 128, 28, 28)
Pool3	(…, 128, 14, 14)
Conv4	(…, 128, 12, 12)
Pool4	(…, 128, 6, 6)
Flatten	(…, 4608)
FC1	(…, 256)
FC2	(…, 1)

Output shape

As a sanity check, let’s pass a batch of 3x128x128 images through the network and check the output shape. Note that the model outputs a 2D tensor of shape (batch_size, 1), where the second dimension is the number of neurons in the output layer, which is 1 in this case.

batch = torch.rand(16, 3, 128, 128)

model = BaselineModel()
output = model(batch)

print("Input shape:", *batch.shape)
print("Output shape:", *output.shape)

Input shape: 16 3 128 128
Output shape: 16 1

Note

The model returns a tensor of shape (batch_size, 1), which is one binary-classification logit for each example. A single output unit does not remove the feature dimension on the second axis.

Training

Now let’s train the baseline model on the cats-vs-dogs dataset. We start by importing the Trainer class from the training.py file. We also import binary accuracy from TorchEval as our evaluation metric.

from training import Trainer
from torcheval.metrics import BinaryAccuracy

Custom adapter

The binary cross-entropy function expects predictions and targets to have matching shapes. In our case, the network returns a prediction tensor of shape (batch_size, 1), while the dataset supplies a target tensor of shape (batch_size,). This will cause a shape mismatch error when computing the loss. In addition, the binary cross-entropy function expects the targets to use a floating-point dtype; otherwise, it will raise a runtime error.

There are several ways to handle these issues.

• One option is to define a custom loss function that removes the singleton dimension from the predictions and converts the targets to floating point. However, this approach can lead to code duplication and inconsistencies if we need to use the same logic in multiple places.

• Another option is to modify the model directly so that it removes the singleton dimension from its output. This is a good solution, but it only fixes the shape mismatch. In future tutorials, we will use it together with a target transformation that converts the labels to floating point directly in the dataset.

• The third option is to define an adapter that prepares the model outputs and targets before passing them to either the loss function or an evaluation metric. This is the approach we will take in this tutorial.

The Trainer class is designed to support different training scenarios. Internally, it uses an adapter function that specifies how a batch is passed through the model and then supplied to a loss function or metric. The default implementation is:

def default_adapter(model, batch, loss_fn):
    inputs, labels = batch
    outputs = model(inputs)
    return loss_fn(outputs, labels)

We can replace this default adapter with a custom one that removes the singleton dimension from the predictions and converts the labels to floating point. The same adapter can then prepare the tensors used by both the loss function and the evaluation metric.

def binary_adapter(model, batch, func):
    inputs, targets = batch
    logits = model(inputs).squeeze(-1)
    return func(logits, targets.float())

Note

We use squeeze(-1) rather than an unrestricted squeeze(). Only the final dimension is removed, while the batch dimension is preserved even when a batch contains a single example.

Training loop

We train the network for several epochs using the binary cross-entropy loss modified to handle the shape mismatch. We use a smaller learning rate of 1e-4 because the model is more complex. We provide the validation data to the trainer so it can monitor the validation loss and accuracy after each epoch. This is useful for recognizing overfitting.

Warning: this cell is computationally expensive without a GPU. (3~4 minutes per epoch)

train_loader = torch.utils.data.DataLoader(train_ds, batch_size=64, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_ds, batch_size=128, shuffle=False)
test_loader = torch.utils.data.DataLoader(test_ds, batch_size=128, shuffle=False)

torch.manual_seed(42)

device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model   = BaselineModel().to(device)
loss_fn = torch.nn.BCEWithLogitsLoss()

optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
epochs = 10

trainer = Trainer()
trainer.set_adapter(binary_adapter)

history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)

===== Training on cuda:0 device =====

Epoch  1/10: 100%|██████████| 274/274 [01:46<00:00,  2.56it/s, train_loss=0.6637, valid_loss=0.6309]
Epoch  2/10: 100%|██████████| 274/274 [00:48<00:00,  5.68it/s, train_loss=0.6118, valid_loss=0.5932]
Epoch  3/10: 100%|██████████| 274/274 [00:47<00:00,  5.77it/s, train_loss=0.5731, valid_loss=0.5560]
Epoch  4/10: 100%|██████████| 274/274 [00:48<00:00,  5.71it/s, train_loss=0.5318, valid_loss=0.5243]
Epoch  5/10: 100%|██████████| 274/274 [00:47<00:00,  5.81it/s, train_loss=0.5046, valid_loss=0.4954]
Epoch  6/10: 100%|██████████| 274/274 [00:47<00:00,  5.76it/s, train_loss=0.4848, valid_loss=0.4815]
Epoch  7/10: 100%|██████████| 274/274 [00:47<00:00,  5.71it/s, train_loss=0.4610, valid_loss=0.4707]
Epoch  8/10: 100%|██████████| 274/274 [00:46<00:00,  5.84it/s, train_loss=0.4467, valid_loss=0.4571]
Epoch  9/10: 100%|██████████| 274/274 [00:48<00:00,  5.64it/s, train_loss=0.4283, valid_loss=0.4393]
Epoch 10/10: 100%|██████████| 274/274 [00:47<00:00,  5.79it/s, train_loss=0.4195, valid_loss=0.4365]

We plot training and validation loss together. The usual pattern is that the training loss decreases over time, while the validation loss decreases initially but then flattens out. The gap between the two curves can indicates overfitting if it widens significantly. In this particular example, the overfitting is minimal, since we stopped training early.

Show code cell source

Hide code cell source

plt.plot(history['train_loss'], label='Train loss')
plt.plot(history['valid_loss'], label='Validation loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

Final evaluation

Once the architecture, hyperparameters, and stopping rule are fixed using training and validation data, we evaluate the final model once on the test set. Accuracy is easy to interpret here because the classes are nearly balanced. (For imbalanced problems, also inspect precision, recall, and the confusion matrix.)

Show code cell source

Hide code cell source

trainer.set_metrics(accuracy=BinaryAccuracy())
ans = trainer.eval(model, test_loader)

print(f"Test Accuracy: {ans['accuracy']:.2%}")

Test Accuracy: 81.76%

Summary

In this tutorial, we learned how to convert a folder of JPEG files into a reproducible image-classification pipeline. We extracted the dataset, validated the images, applied deterministic preprocessing, inspected the class mapping, and created stratified train/validation/test splits. We then trained a small CNN while monitoring validation loss and reserved the test set for final evaluation.