Using PyMC3 on the GPU #1246
Description
Most input provided by @fhuszar.
It seems there are at least two blockers of using pymc3 on the gpu. The first one is incompatilibty with float32 dtype. Here is an example model:
from pymc3 import Model, NUTS, sample
from pymc3.distributions import DensityDist
import pymc3 as pm
import theano
import numpy as np
theano.config.floatX = 'float32'
theano.config.compute_test_value = 'raise'
theano.config.exception_verbosity= 'high'
with Model() as denoising_model:
# theano.config.compute_test_value = 'off'
x = DensityDist('x',
logp= lambda value: -(value**2).sum(),
shape=(1, 1, 10, 10),
testval=np.random.randn(1,1,10,10).astype('float32'),
dtype='float32',
)
sampler = pm.Metropolis()
trace = sample(10, sampler)Output:
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
/home/wiecki/miniconda3/lib/python3.5/site-packages/theano/compile/function_module.py in __call__(self, *args, **kwargs)
898 outputs =\
--> 899 self.fn() if output_subset is None else\
900 self.fn(output_subset=output_subset)
TypeError: expected type_num 11 (NPY_FLOAT32) got 12
During handling of the above exception, another exception occurred:
TypeError Traceback (most recent call last)
<ipython-input-1-9a14b0f25285> in <module>()
19
20 sampler = pm.Metropolis()
---> 21 trace = sample(10, sampler)
/home/wiecki/working/projects/pymc/pymc3/sampling.py in sample(draws, step, start, trace, chain, njobs, tune, progressbar, model, random_seed)
148 sample_func = _sample
149
--> 150 return sample_func(**sample_args)
151
152
/home/wiecki/working/projects/pymc/pymc3/sampling.py in _sample(draws, step, start, trace, chain, tune, progressbar, model, random_seed)
157 progress = progress_bar(draws)
158 try:
--> 159 for i, strace in enumerate(sampling):
160 if progressbar:
161 progress.update(i)
/home/wiecki/working/projects/pymc/pymc3/sampling.py in _iter_sample(draws, step, start, trace, chain, tune, model, random_seed)
239 if i == tune:
240 step = stop_tuning(step)
--> 241 point = step.step(point)
242 strace.record(point)
243 yield strace
/home/wiecki/working/projects/pymc/pymc3/step_methods/arraystep.py in step(self, point)
125 bij = DictToArrayBijection(self.ordering, point)
126
--> 127 apoint = self.astep(bij.map(point))
128 return bij.rmap(apoint)
129
/home/wiecki/working/projects/pymc/pymc3/step_methods/metropolis.py in astep(self, q0)
125 q = q0 + delta
126
--> 127 q_new = metrop_select(self.delta_logp(q, q0), q, q0)
128
129 if q_new is q:
/home/wiecki/miniconda3/lib/python3.5/site-packages/theano/compile/function_module.py in __call__(self, *args, **kwargs)
910 node=self.fn.nodes[self.fn.position_of_error],
911 thunk=thunk,
--> 912 storage_map=getattr(self.fn, 'storage_map', None))
913 else:
914 # old-style linkers raise their own exceptions
/home/wiecki/miniconda3/lib/python3.5/site-packages/theano/gof/link.py in raise_with_op(node, thunk, exc_info, storage_map)
312 # extra long error message in that case.
313 pass
--> 314 reraise(exc_type, exc_value, exc_trace)
315
316
/home/wiecki/miniconda3/lib/python3.5/site-packages/six.py in reraise(tp, value, tb)
683 value = tp()
684 if value.__traceback__ is not tb:
--> 685 raise value.with_traceback(tb)
686 raise value
687
/home/wiecki/miniconda3/lib/python3.5/site-packages/theano/compile/function_module.py in __call__(self, *args, **kwargs)
897 try:
898 outputs =\
--> 899 self.fn() if output_subset is None else\
900 self.fn(output_subset=output_subset)
901 except Exception:
TypeError: expected type_num 11 (NPY_FLOAT32) got 12
Apply node that caused the error: Elemwise{sqr,no_inplace}(Reshape{4}.0)
Toposort index: 5
Inputs types: [TensorType(float32, (True, True, False, False))]
Inputs shapes: [(1, 1, 10, 10)]
Inputs strides: [(800, 800, 80, 8)]
Inputs values: ['not shown']
Outputs clients: [[Sum{acc_dtype=float64}(Elemwise{sqr,no_inplace}.0)]]
Debugprint of the apply node:
Elemwise{sqr,no_inplace} [id A] <TensorType(float32, (True, True, False, False))> ''
|Reshape{4} [id B] <TensorType(float32, (True, True, False, False))> ''
|Subtensor{int64:int64:} [id C] <TensorType(float32, vector)> ''
| |inarray1 [id D] <TensorType(float32, vector)>
| |Constant{0} [id E] <int64>
| |Constant{100} [id F] <int64>
|TensorConstant{[ 1 1 10 10]} [id G] <TensorType(int64, vector)>
Storage map footprint:
- inarray, Input, Shape: (100,), ElemSize: 8 Byte(s), TotalSize: 800 Byte(s)
- Reshape{4}.0, Shape: (1, 1, 10, 10), ElemSize: 8 Byte(s), TotalSize: 800 Byte(s)
- inarray1, Input, Shape: (100,), ElemSize: 8 Byte(s), TotalSize: 800 Byte(s)
- TensorConstant{[ 1 1 10 10]}, Shape: (4,), ElemSize: 8 Byte(s), TotalSize: 32 Byte(s)
- Constant{0}, Shape: (), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
- Constant{100}, Shape: (), ElemSize: 8 Byte(s), TotalSize: 8.0 Byte(s)
TotalSize: 2448.0 Byte(s) 0.000 GB
TotalSize inputs: 1648.0 Byte(s) 0.000 GB
HINT: Re-running with most Theano optimization disabled could give you a back-trace of when this node was created. This can be done with by setting the Theano flag 'optimizer=fast_compile'. If that does not work, Theano optimizations can be disabled with 'optimizer=None'.
Not sure where the float64 dtype comes in. @nouiz any ideas?
The second problem is related to #566 and should have a simple solution of only setting the test value behavior in the model context.