matterport / Mask_RCNN

Mask R-CNN for object detection and instance segmentation on Keras and TensorFlow
Other
24.73k stars 11.71k forks source link

Run Mask RCNN on GPU instead of CPU #2464

Open Basma-hash opened 3 years ago

Basma-hash commented 3 years ago

Hello Everyone, I would like to run my code on GPU instead of CPU anyone can help me please?

RayVik commented 3 years ago

@Basma-hash, part my code for run on GPU:

import os
import warnings

import mrcnn.model as modellib
import tensorflow as tf

from keras import backend as K
from mrcnn.config import Config

os.environ["CUDA_VISIBLE_DEVICES"] = '0'
gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=per_process_gpu_memory)
cfg = tf.compat.v1.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True)  
sess = tf.compat.v1.Session(config=cfg)
K.set_session(sess)

class CustomConfig(Config):
      GPU_COUNT = 1
      IMAGES_PER_GPU = 1
      NUM_CLASSES = 1 + 1
      NAME = "balloon"
      STEPS_PER_EPOCH = 100
      DETECTION_MIN_CONFIDENCE = 0.9

with tf.device('/job:localhost/replica:0/task:0/device:GPU:0'):
    rcnn = modellib.MaskRCNN(mode='inference', model_dir=MODEL_DIR,
                             config=CustomConfig())

    rcnn.load_weights(MODEL_DIR, by_name=True)
    graph = tf.get_default_graph()

I hope is help

Malandru commented 3 years ago

Check your tensorflow version with pip freeze. Then, uninstall tensorflow and install tensorflow-gpu.

For example:

pip uninstall tensorflow
pip install tensorflow-gpu==x.y.z

Replace x.y.z with the version that was printed in pip freeze for tensorflow. I recommend you to do this in a virtual environment.