This project trains a deep learning image classification model using MobileNetV2 with PyTorch. It supports custom datasets, performs data augmentation, and includes a training + validation pipeline with progress tracking.
The trained model is saved along with the class labels for future inference.
Uses Transfer Learning (MobileNetV2 pretrained on ImageNet)
Supports custom datasets
Automatic GPU usage if CUDA is available
Data augmentation for better generalization
Training + validation accuracy tracking
Progress bars with tqdm
Saves model weights + class labels
Install the required libraries:
pip install torch torchvision tqdm
Optional (recommended for faster loading):
pip install accelerate
Before running the script, update these paths in the code:
TRAIN_DIR = "path_to_training_data" VAL_DIR = "path_to_validation_data" MODEL_PATH = "model.pth"
Other parameters:
BATCH_SIZE = 16 EPOCHS = 20 LR = 0.001
The following augmentations are applied to training images:
Resize to 224x224
Random horizontal flip
Random rotation
Color jitter (brightness, contrast, saturation)
This improves model generalization.
#🏋️ Training the Model
Run the training script:
python train_model.py
During training you will see progress like:
Epoch 1/20 [Train]: 100%|██████████| Train Loss: 0.5432 | Train Acc: 0.8123 Val Loss: 0.4312 | Val Acc: 0.8541
After training completes, the model is saved as:
model.pth
It contains:
Model weights
Class names
Example saved format:
{ "model_state_dict": ..., "class_names": [...] }
The script automatically selects:
GPU (CUDA) if available
CPU otherwise
Using device: cuda
or
Using device: cpu
The project uses:
MobileNetV2
Advantages:
Lightweight
Fast training
Good performance on small datasets
The final classifier layer is modified to match the number of dataset classes.
The model loads pretrained weights:
ImageNet pretrained weights
Then only the final layer is adapted for the custom dataset.
This significantly improves performance on smaller datasets.