How can I prevent overfitting in a convolutional neural network while training on a small?

import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms # Dataset train_dataset = datasets.ImageFolder('data/train', transform=transforms.ToTensor()) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True) # Simple CNN class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(16*32*32, 10) def forward(self, x): x = self.pool(torch.relu(self.conv1(x))) x = x.view(-1, 16*32*32) x = self.fc1(x) return x model = SimpleCNN() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # Training loop for epoch in range(10): for inputs, labels in train_loader: optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step()

Training a CNN on a small dataset almost always increases the risk of overfitting because the model starts memorizing the training images instead of learning features that generalize well. The goal is to reduce model complexity while increasing the diversity of the data the model sees during training. Common techniques include:

Data Augmentation (most effective for small datasets)
- Random horizontal/vertical flips
- Random rotations
- Random crops
- Brightness and contrast adjustments
- Zoom and translation
Transfer Learning
- Instead of training a CNN from scratch, start with a pretrained model like ResNet50, EfficientNet, or MobileNet.
- Freeze the backbone initially and only train the classifier head.
- Fine-tune later using a low learning rate.
Regularization
- Add Dropout (0.3–0.5)
- Use L2 (weight decay)
- Apply Batch Normalization
Early Stopping
- Stop training once the validation loss stops improving and restore the best model.
Reduce Model Complexity
- If you’re using a deep custom CNN, try fewer convolutional filters or layers.
- A simpler model often performs better on limited data.
Cross Validation
- K-fold cross-validation provides a more reliable estimate of model performance when the dataset is small.

Example using TensorFlow/Keras:

import tensorflow as tf
from tensorflow.keras import layers, models, regularizers
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Data Augmentation
train_datagen = ImageDataGenerator(
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest'
)

model = models.Sequential([
    layers.Conv2D(32, (3,3), activation='relu', input_shape=(224,224,3)),
    layers.BatchNormalization(),
    layers.MaxPooling2D(),

    layers.Conv2D(
        64, (3,3),
        activation='relu',
        kernel_regularizer=regularizers.l2(1e-4)
    ),
    layers.BatchNormalization(),
    layers.MaxPooling2D(),

    layers.Flatten(),
    layers.Dropout(0.5),
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

early_stop = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',
    patience=5,
    restore_best_weights=True
)

history = model.fit(
    train_generator,
    validation_data=val_generator,
    epochs=50,
    callbacks=[early_stop]
)

For very small datasets (a few thousand images or fewer), transfer learning with a pretrained model plus strong data augmentation usually produces significantly better results than training a CNN from scratch. If possible, also collect additional training data, as more high-quality data generally provides the biggest improvement in generalization.