Detectron2 Environment Setup and Installation on a Local Computer [GPU Enabled]

Before I write about setting up Detectron2 on a local computer with a GPU, let me tell you something about it.
Detectron2 is the cutting-edge library developed by Facebook AI Research (FAIR) which provides detection and segmentation algorithms. It started with Detectron in 2018, which was implemented in Caffe2. By incorporating state-of-the-art neural network architectures, Detectron2 ensures accuracy and efficiency in detecting and segmenting objects. Its comprehensive functionality allows for a wide range of applications, making it an indispensable resource for the computer vision community. Its capabilities include panoptic segmentation, Densepose, Cascade R-CNN, rotated bounding boxes, PointRend, DeepLab, ViTDet, MViTv2 etc. The models can be exported to TorchScript format or Caffe2 format for deployment. And, it trains faster! You can learn more about it on this blog.
The Detectron2 platform is implemented using CUDA and PyTorch. It can provide fast training even with a single GPU. It has a modular and extensible design that allows users to plug in custom modules into almost any part of the object detection system. The dependencies required for Detectron 2 are Python ≥ 3.6 and PyTorch ≥ 1.6. If your system does not have a GPU, you can use your CPU for training and inferencing, but you need to be patient about the speed of training. I tried training with the CPU on my other laptop, and it required overnight to complete the training. But with GPU, it was trained within an hour.
Detectron2 benefits from ongoing updates and improvements, as Facebook AI Research remains committed to pushing the boundaries of AI technology. This commitment ensures that the library remains at the forefront of the field, constantly evolving to meet the demands of emerging research and real-world applications.
I had been trying to install Detectron 2 on my Windows 11 (with NVIDIA RTX A4500) for a month, and I figured out that it was the CUDA problem and nothing had to do with other dependencies.
Following are the steps that worked for me:
- I followed a YouTube tutorial to uninstall the existing CUDA Toolkit. You need to uninstall multiple programs and be very careful not to remove important ones.
- You’ll need to install the CUDA Toolkit along with a version compatible with the PyTorch version. I wanted to install Detectron2 0.6 with Torch 1.8.0, so I installed CUDA Toolkit 11.1. After installing CUDA and downloading the cuDNN compatible for the cuda version installed, you’ll need to provide the cuDNN path to the system environment variables. For that, you can follow this tutorial. You need to reboot after making changes to the system environment variables.
- Once you have fixed CUDA and cuDNN, create a new virtual environment in Anaconda prompt: conda create --new dvenv python=3.8
- Activate the environment: conda activate dvenv
- Install pyyaml: pip install pyyaml==5.1
- For CUDA 11.1, install torch, torchvision, and torchaudio using this command: pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html For other versions of CUDA or the CPU, you can check out previous versions of PyTorch from here.
- Install fvcore: pip install git+https://github.com/facebookresearch/fvcore
- Install cython: pip install cython
- Install OpenCv-python: pip install opencv-python
- Install pycocotools: pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI
- Now, to install Detectron2, you need to clone the official repo from Facebook Research: git clone https://github.com/facebookresearch/detectron2.git
- Get inside the detectron2 folder that you have just cloned: cd detectron2
- Run: python setup.py build develop
- Run: python train.py
After the 11th step, Detectron 2 was installed in my virtual environment. But when I ran python train.py, I got an error that said ModuleNotFoundError: No module named google.protobuf
To solve this error, I uninstalled protobuf and reinstalled it with pip. Though it said the protobuf was installed successfully, I got the same error when running train.py. I fixed it by installing protobuf using Conda. Strange, right? Anyway, it runs after this! You can train your own model using custom datasets with detectron2 now!
My github has the code and other instructions for training Detectron2 with custom datasets for a single class. You can find it here.
In my next article, I’ll write about preparing suitable datasets for Detectron2 and what steps I used to train it on my custom dataset. You can check out my article on creating annotations on your images using makesense.ai, but make sure you export it in COCO format at the end. Moreover, I'll explore multiple capabilities of Detectron2 like panoptic segmentation, densepose, cascade R-CNN, etc., and steps to do them.