Subscribe to DSC Newsletter

Machine Learning with C++ - Mask R-CNN with PyTorch C++ Frontend

I made C++ implementation of Mask R-CNN with PyTorch C++ frontend. The code is based on PyTorch implementations from multimodallearning and Keras implementation from Matterport . Project was made for educational purposes and can be used as comprehensive example of PyTorch C++ frontend API. Besides regular API you will find how to: load data from MSCoco dataset, create custom layers, manage weights across language boundaries(Python to C++).

This implementation use heads from resnet50 and can be trained only with coco dataset.

I made development on custom PyTorch build based on this revision, because at that moment there was no Cuda 10 support in published binaries. Now I assume you can use binaries for PyTorch v1.x.

Development environment configuration

I'm using Arch Linux, with additional packages openblas, OpenCV, gcc-7, cuda. At the moment when I was building PyTorch Cuda had support only for gcc-7 as host compiler, so you need to configure a build to use it. Don't use CC environmental variable for compiler configuration, because scripts depend on gcc. To make build successful I used next strategy: created a directory $HOME/old_gcc, then symlink gcc in that directory to /usr/bin/gcc-7. Add that directory to the front of PATH export PATH=$HOME/old_gcc:$PATH before building PyTorch.

Install PyYaml for your python environment pip install pyyaml.

PyTorch python wheel compilation

This step can be skipped if you don't need python environment with same version of PyTorch

PyTorch has scripts for building library from sources, but before run them you need to setup environment variables, I used next ones:

  • MAX_JOBS=8
  • ONNX_NAMESPACE=onnx_torch

And used next commands to build a wheel:

python bdist_wheel 
pip install torch-1.0.0a0+4f0434d-cp37-cp37m-linux_x86_64.whl

PyTorch C++ Frontend Compilation

If don't need a python wheel for PyTorch you can build only a C++ part. The previous step also builds the C++ frontend.

PyTorch has a CMake scripts, which can be used for build configuration and compilation. So you can use general procedure for building projects with CMake. I used next CMake command-line parameters to be able to build PyTorch in my environment:

  • USE_CUDA=1

I changed CMake parameter CMAKE_PREFIX_PATH to use custom directory for PyTorch installation.

Application compilation

  1. After checking out of the code please update also submodules for the project, it have dependency for Eigen and Json parser libraries.
  2. Update CMake parameter CMAKE_PREFIX_PATH with path where you installed PyTorch libraries, it will make find_package(Torch REQUIRED) works.

Parameters management

Please notice that parameters saved from python version of PyTorch with save_state_dict function are saved with pickle module, so are incompatible with C++ loading routings from PyTorch C++ frontend. How to manage parameters across language boundaries see code and comments in sateloader.h file.


There are two projects mask-rcnn_demo and mask-rcnn_train which should be used with next parameters:

  • Demo - mask-rcnn_demo executable takes two parameters path to file with trained parameters and path to image file for classification. You can use pre-trained parameters from the original project (I just converted them to the format acceptable for C++ application). After processing you will get file, named result.png in your's working directory, with rendered bounding boxes, masks and printed labels. Command line can looks like this "mask-rcnn_demo test.png"

  • Train - mask-rcnn_train executable takes twp parameters path to the coco dataset and path to the pretrained model. If you want to start training from scratch, please put path to the pretrained resnet50 weights. Command line can looks like this "mask-rcnn_train /development/data/coco /development/model/". Default name for check-point file is ./logs/ You can download pre-trained resnet parameters from here.

Also you can download file with pre-trained parameters from this link, it was made for proof of the concept and for vehicles label types only also it was trained on small number of iteration, because I don't have suitable hardware for full training cycle.



Views: 2576

Tags: C++, cnn, deep learning, image, machine learning, mask


You need to be a member of Data Science Central to add comments!

Join Data Science Central


  • Add Videos
  • View All

© 2020   Data Science Central ®   Powered by

Badges  |  Report an Issue  |  Privacy Policy  |  Terms of Service