tensorflow / tensorflow/probability

MC sampling with GPU does not work

Open
#483 0 comments 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Jupyter Notebook
Stars
4.4k
Forks
1.1k
No merged PRs in 30d

Description

Hi All,

I would like to use "tfp.mcmc.sample_chain" for sampling a model, yet unfortunately it does not work for GPU device. Can you please look into it? Thanks!

This is the used code:

`

import tensorflow as tf
import numpy as np
import sys
import tensorflow_probability as tfp
tfd = tfp.distributions


with tf.device('/gpu:0'):
    chain_num = 100
    data_dims = [chain_num, 2]
    init_samples = tf.random_uniform(data_dims, 0.0, 1.0, tf.float32)
    step_size=0.1
    num_steps = 20

    def build_f_func(x):
        f = tf.reduce_sum(tf.square(x), axis = -1, keepdims=True)
        return tf.squeeze(f, axis = -1)

    hmc = tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=build_f_func,
        num_leapfrog_steps=3,
        step_size=step_size)
    
with tf.device('/cpu:0'):
    samples, is_accepted = tfp.mcmc.sample_chain(
         num_results=1,
         current_state=init_samples,
         kernel=hmc,
         num_burnin_steps=num_steps,
         trace_fn=lambda _, pkr: pkr.is_accepted)
    
with tf.device('/gpu:0'):
    samples = tf.stop_gradient(samples)
    samples = tf.squeeze(samples, axis = 0)
    
    is_accepted = tf.reduce_mean(tf.cast(is_accepted, dtype=tf.float32))
    
    
init_op = tf.compat.v1.global_variables_initializer()

sess = tf.Session()
sess.run(init_op)

`

If second "tf.device('/cpu:0')" is replaced with "tf.device('/gpu:0')", I get error:
`

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation mcmc_sample_chain_1/one_hot: Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=-1 requested_device_name_='/device:GPU:0' assigned_device_name_='' resource_device_name_='' supported_device_types_=[CPU, XLA_CPU, XLA_GPU] possible_devices_=[]
TensorArrayScatterV3: CPU XLA_CPU XLA_GPU 
TensorArrayReadV3: CPU XLA_CPU XLA_GPU 
Enter: GPU CPU XLA_CPU XLA_GPU 
TensorArrayV3: CPU XLA_CPU XLA_GPU 
OneHot: GPU CPU XLA_CPU XLA_GPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  mcmc_sample_chain_1/one_hot (OneHot) /device:GPU:0
  mcmc_sample_chain_1/trace_scan/TensorArray (TensorArrayV3) 
  mcmc_sample_chain_1/trace_scan/TensorArrayUnstack/TensorArrayScatter/TensorArrayScatterV3 (TensorArrayScatterV3) /device:GPU:0
  mcmc_sample_chain_1/trace_scan/while/TensorArrayReadV3/Enter (Enter) /device:GPU:0
  mcmc_sample_chain_1/trace_scan/while/TensorArrayReadV3 (TensorArrayReadV3) /device:GPU:0

Op: OneHot
Node attrs: T=DT_INT32, TI=DT_INT64, axis=-1
Registered kernels:
  device='XLA_CPU'; TI in [DT_INT32, DT_UINT8, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT8, ..., DT_BFLOAT16, DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64]
  device='XLA_GPU'; TI in [DT_INT32, DT_UINT8, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT8, ..., DT_QINT32, DT_BFLOAT16, DT_HALF, DT_UINT32, DT_UINT64]
  device='XLA_CPU_JIT'; TI in [DT_INT32, DT_UINT8, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT8, ..., DT_BFLOAT16, DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64]
  device='XLA_GPU_JIT'; TI in [DT_INT32, DT_UINT8, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT8, ..., DT_QINT32, DT_BFLOAT16, DT_HALF, DT_UINT32, DT_UINT64]
  device='CPU'; TI in [DT_INT64]; T in [DT_VARIANT]
  device='CPU'; TI in [DT_INT32]; T in [DT_VARIANT]
  device='CPU'; TI in [DT_UINT8]; T in [DT_VARIANT]
  device='CPU'; TI in [DT_INT64]; T in [DT_RESOURCE]
  device='CPU'; TI in [DT_INT32]; T in [DT_RESOURCE]
  device='CPU'; TI in [DT_UINT8]; T in [DT_RESOURCE]
  device='CPU'; TI in [DT_INT64]; T in [DT_STRING]
  device='CPU'; TI in [DT_INT32]; T in [DT_STRING]
  device='CPU'; TI in [DT_UINT8]; T in [DT_STRING]
  device='CPU'; TI in [DT_INT64]; T in [DT_BOOL]
  device='CPU'; TI in [DT_INT32]; T in [DT_BOOL]
  device='CPU'; TI in [DT_UINT8]; T in [DT_BOOL]
  device='CPU'; TI in [DT_INT64]; T in [DT_COMPLEX128]
  device='CPU'; TI in [DT_INT32]; T in [DT_COMPLEX128]
  device='CPU'; TI in [DT_UINT8]; T in [DT_COMPLEX128]
  device='CPU'; TI in [DT_INT64]; T in [DT_COMPLEX64]
  device='CPU'; TI in [DT_INT32]; T in [DT_COMPLEX64]
  device='CPU'; TI in [DT_UINT8]; T in [DT_COMPLEX64]
  device='CPU'; TI in [DT_INT64]; T in [DT_DOUBLE]
  device='CPU'; TI in [DT_INT32]; T in [DT_DOUBLE]
  device='CPU'; TI in [DT_UINT8]; T in [DT_DOUBLE]
  device='CPU'; TI in [DT_INT64]; T in [DT_FLOAT]
  device='CPU'; TI in [DT_INT32]; T in [DT_FLOAT]
  device='CPU'; TI in [DT_UINT8]; T in [DT_FLOAT]
  device='CPU'; TI in [DT_INT64]; T in [DT_BFLOAT16]
  device='CPU'; TI in [DT_INT32]; T in [DT_BFLOAT16]
  device='CPU'; TI in [DT_UINT8]; T in [DT_BFLOAT16]
  device='CPU'; TI in [DT_INT64]; T in [DT_HALF]
  device='CPU'; TI in [DT_INT32]; T in [DT_HALF]
  device='CPU'; TI in [DT_UINT8]; T in [DT_HALF]
  device='CPU'; TI in [DT_INT64]; T in [DT_INT8]
  device='CPU'; TI in [DT_INT32]; T in [DT_INT8]
  device='CPU'; TI in [DT_UINT8]; T in [DT_INT8]
  device='CPU'; TI in [DT_INT64]; T in [DT_UINT8]
  device='CPU'; TI in [DT_INT32]; T in [DT_UINT8]
  device='CPU'; TI in [DT_UINT8]; T in [DT_UINT8]
  device='CPU'; TI in [DT_INT64]; T in [DT_INT16]
  device='CPU'; TI in [DT_INT32]; T in [DT_INT16]
  device='CPU'; TI in [DT_UINT8]; T in [DT_INT16]
  device='CPU'; TI in [DT_INT64]; T in [DT_UINT16]
  device='CPU'; TI in [DT_INT32]; T in [DT_UINT16]
  device='CPU'; TI in [DT_UINT8]; T in [DT_UINT16]
  device='CPU'; TI in [DT_INT64]; T in [DT_INT32]
  device='CPU'; TI in [DT_INT32]; T in [DT_INT32]
  device='CPU'; TI in [DT_UINT8]; T in [DT_INT32]
  device='CPU'; TI in [DT_INT64]; T in [DT_INT64]
  device='CPU'; TI in [DT_INT32]; T in [DT_INT64]
  device='CPU'; TI in [DT_UINT8]; T in [DT_INT64]
  device='GPU'; TI in [DT_INT64]; T in [DT_INT64]
  device='GPU'; TI in [DT_INT32]; T in [DT_INT64]
  device='GPU'; TI in [DT_UINT8]; T in [DT_INT64]
  device='GPU'; TI in [DT_INT64]; T in [DT_INT32]
  device='GPU'; TI in [DT_INT32]; T in [DT_INT32]
  device='GPU'; TI in [DT_UINT8]; T in [DT_INT32]
  device='GPU'; TI in [DT_INT64]; T in [DT_BOOL]
  device='GPU'; TI in [DT_INT32]; T in [DT_BOOL]
  device='GPU'; TI in [DT_UINT8]; T in [DT_BOOL]
  device='GPU'; TI in [DT_INT64]; T in [DT_DOUBLE]
  device='GPU'; TI in [DT_INT32]; T in [DT_DOUBLE]
  device='GPU'; TI in [DT_UINT8]; T in [DT_DOUBLE]
  device='GPU'; TI in [DT_INT64]; T in [DT_FLOAT]
  device='GPU'; TI in [DT_INT32]; T in [DT_FLOAT]
  device='GPU'; TI in [DT_UINT8]; T in [DT_FLOAT]
  device='GPU'; TI in [DT_INT64]; T in [DT_HALF]
  device='GPU'; TI in [DT_INT32]; T in [DT_HALF]
  device='GPU'; TI in [DT_UINT8]; T in [DT_HALF]

     [[node mcmc_sample_chain_1/one_hot (defined at usr/local/lib/python2.7/dist-packages/tensorflow_probability/python/mcmc/sample.py:356) ]]

Errors may have originated from an input operation.
Input Source operations connected to node mcmc_sample_chain_1/one_hot:
 mcmc_sample_chain_1/num_results (defined at usr/local/lib/python2.7/dist-packages/tensorflow_probability/python/mcmc/sample.py:315)    
 mcmc_sample_chain_1/add_1 (defined at usr/local/lib/python2.7/dist-packages/tensorflow_probability/python/mcmc/sample.py:355)    
 mcmc_sample_chain_1/add (defined at usr/local/lib/python2.7/dist-packages/tensorflow_probability/python/mcmc/sample.py:354)

Original stack trace for u'mcmc_sample_chain_1/one_hot':
  File "/draw_points_along_sampling_tf.py", line 192, in <module>
    samples = sample_points()
  File "/draw_points_along_sampling_tf.py", line 146, in sample_points
    trace_fn=lambda _, pkr: pkr.is_accepted)
  File "usr/local/lib/python2.7/dist-packages/tensorflow_probability/python/mcmc/sample.py", line 356, in sample_chain
    dtype=tf.int32),
  File "usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/array_ops.py", line 3027, in one_hot
    name)
  File "usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/gen_array_ops.py", line 5731, in one_hot
    off_value=off_value, axis=axis, name=name)
  File "usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "usr/local/lib/python2.7/dist-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/ops.py", line 3616, in create_op
    op_def=op_def)
  File "usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/ops.py", line 2005, in __init__
    self._traceback = tf_stack.extract_stack()

`

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at tensorflow_probability/python/mcmc/sample.py:356, the sample_chain entry point named in the traceback, and reproduce the provided GPU configuration. Investigate the mcmc_sample_chain_1/one_hot device error and verify that GPU execution of sample_chain completes without the reported InvalidArgumentError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.