Framework for training a neural network model for object detection in 2d image
Make sure your dataset folder is in following order.
dataset
├── train
│ ├── images # raw images
│ ├── labels # bbox label files
└── test
│ ├── images # raw images
│ ├── labels # bbox label files
└── valid
│ ├── images # raw images
│ ├── labels # bbox label files
└── models # directory to save the checkpoints
│ ├── ... # saved checkpoints
└── tensorboard_logs # directory to save tensorboard logs
├── ... # saved logs
Please create
models
andtensorboard_logs
directories if they don't exist.
Make sure your image file name and label file name is same i.e., if image file name issome_image.png
then corresponding label file would besome_image.txt
Distributed
python -m torch.distributed.launch --nproc_per_node=<number_of_avaliable_GPUs> train.py
--data <path_to_dataset_folder> --batch-size <N int> --log-dir <tensorboard_logs_folder> --save-as <filename>
Single GPU
python train.py
--data <path_to_dataset_folder> --batch-size <N int> --log-dir <tensorboard_logs_folder> --save-as <filename>