提交 b8e3cf3f authored 作者: Frederic's avatar Frederic

Change how we set the pycuda device to be sure we always have the good one initialized by Theano.

上级 32d12f3c
import os import os
import warnings
import theano import theano
import theano.sandbox.cuda as cuda import theano.sandbox.cuda as cuda
cuda_ndarray = cuda.cuda_ndarray.cuda_ndarray
def select_gpu_from_theano(): def set_gpu_from_theano():
# Transfer the theano gpu binding to pycuda, for consistency """
theano_to_pycuda_device_map = {"cpu": "0", This set the GPU used by PyCUDA to the same as the one used by Theano.
"gpu0": "0", """
"gpu1": "1", #import pdb;pdb.set_trace()
"gpu2": "2", if cuda.use.device_number is None:
"gpu3": "3"} cuda.use("gpu",
dev = theano_to_pycuda_device_map.get(theano.config.device, "0") force=False,
if theano.config.device == 'gpu': default_to_move_computation_to_gpu=False,
dev = str(cuda.cuda_ndarray.cuda_ndarray.active_device_number()) move_shared_float32_to_gpu=False,
os.environ["CUDA_DEVICE"] = dev enable_cuda=True,
test_driver=True)
select_gpu_from_theano() assert cuda.use.device_number == cuda_ndarray.active_device_number()
# os.environ["CUDA_DEVICE"] = str(cuda.use.device_number)
set_gpu_from_theano()
pycuda_available = False pycuda_available = False
try: if True: # theano.sandbox.cuda.use.device_number is None:
import pycuda try:
import pycuda.autoinit import pycuda
pycuda_available = True import pycuda.autoinit
except ImportError: pycuda_available = True
# presumably, the user wanted to use pycuda, else they wouldn't have except ImportError:
# imported this module, so issue a warning that the import failed. # presumably, the user wanted to use pycuda, else they wouldn't have
import warnings # imported this module, so issue a warning that the import failed.
warnings.warn("PyCUDA import failed in theano.misc.pycuda_init") warnings.warn("PyCUDA import failed in theano.misc.pycuda_init")
else:
warnings.warn("theano.misc.pycuda_init must be imported before theano"
" init its GPU")
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论