Image Classification with CNN

This project trains a simple Convolutional Neural Network (CNN) on the CIFAR-10 dataset for image classification into 10 categories: plane, car, bird, cat, deer, dog, frog, horse, ship, truck.

Setup

1. Install dependencies:

   pip install -r requirements.txt
   

2. GPU recommended for faster training.

Training

Run the training script:

python train.py

This downloads CIFAR-10, trains the CNN for 5 epochs, and saves the model to ./cifar10_cnn.pth. It also prints test accuracy.

Note: Training takes time; adjust num_epochs for more training.

Inference

After training, run:

python predict.py

This loads the model and predicts on a random test image, showing predicted vs actual class.

Troubleshooting

• CUDA out of memory: Reduce batch size in DataLoader.
• Slow training: Use GPU; CIFAR-10 is small, so CPU works but slower.
• Import errors: Ensure PyTorch is installed correctly for your system.

Requirements

• Python 3.8+
• PyTorch with CUDA support for GPU