提交 4e694699 authored 作者: Frederic Bastien's avatar Frederic Bastien

Fix comment. Fix the force option. Add docstring.

上级 9321a170
......@@ -114,15 +114,39 @@ init_dev.devmap = {}
def use(device,
force=False,
default_to_move_computation_to_gpu=True,
move_shared_float32_to_gpu=True):
move_shared_to_gpu=True):
"""
Error and warning about CUDA should be displayed only when this
function is called. We need to be able to load this module only
to check if it is available!
Parameters
----------
device : string
"cuda", "cuda0", "cudaN", "" (N is the device number to use).
"" mean do all the rest and don't init a device.
force
Will always raise an exception if we can't use the gpu.
default_to_move_computation_to_gpu
If gpu init succeeded, enable by default optimizations to move
computations to the gpu.
move_shared_to_gpu
If gpu init succeeded, put new shared variables on the gpu.
"""
if force:
if not device.startswith('cuda'):
raise Exception("forced the init and bad device provided: " +
device)
else:
# If we force, the device should not already be initialized.
assert device not in init_dev.devmap
if device:
init_dev(device)
elif forece:
raise Exception("No device provided")
if default_to_move_computation_to_gpu:
optdb.add_tags('gpuarray_opt', 'fast_run', 'fast_compile')
optdb.add_tags('gpua_scanOp_make_inplace', 'fast_run')
if move_shared_float32_to_gpu:
if move_shared_to_gpu:
import theano.compile
theano.compile.shared_constructor(gpuarray_shared_constructor)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论