Commit 61af56d9 authored by Yuxin Wu's avatar Yuxin Wu

[MaskRCNN] allow inference with MKL support (#1024)

parent 99565998
...@@ -421,7 +421,9 @@ if __name__ == '__main__': ...@@ -421,7 +421,9 @@ if __name__ == '__main__':
DetectionDataset() # initialize the config with information from our dataset DetectionDataset() # initialize the config with information from our dataset
if args.visualize or args.evaluate or args.predict: if args.visualize or args.evaluate or args.predict:
assert tf.test.is_gpu_available() if not tf.test.is_gpu_available():
from tensorflow.python.framework import test_util
assert test_util.IsMklEnabled(), "Inference requires either GPU support or MKL support!"
assert args.load assert args.load
finalize_configs(is_training=False) finalize_configs(is_training=False)
......
...@@ -44,7 +44,7 @@ def get_num_gpu(): ...@@ -44,7 +44,7 @@ def get_num_gpu():
return ret return ret
env = os.environ.get('CUDA_VISIBLE_DEVICES', None) env = os.environ.get('CUDA_VISIBLE_DEVICES', None)
if env is not None: if env:
return warn_return(len(env.split(',')), "Found non-empty CUDA_VISIBLE_DEVICES. ") return warn_return(len(env.split(',')), "Found non-empty CUDA_VISIBLE_DEVICES. ")
output, code = subproc_call("nvidia-smi -L", timeout=5) output, code = subproc_call("nvidia-smi -L", timeout=5)
if code == 0: if code == 0:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment