Image-Classification-Framework

This is a simple and useful image classification framework based on pytorch. You can easily modify it for other deep learning tasks.

Requirements

This code are tested under python 3.8.0, torch 1.8.0 and torchvision 0.9.0. Other dependencies are listed in requirements.txt. You can install them with the following commands:

pip install -r requirements.txt

Datasets

Dataset link: https://www.kaggle.com/datasets/googolples/cat-dog-images

If you want to run this code on custom dataset, you just need to set the root of your data in --data_root and organize your data like the example dataset.

Run

python main.py