No GPU support for graph convolutions

[IDSPAGET]

Hi,

I have the same issue referenced at the end of https://github.com/tensorflow/graphics/issues/13 Namely the ops SparseReshape does not seem to have GPU support. My code runs fine with automatic device placement but gives this when I force the placement of graph convolutions on gpu0.

[...]
InvalidArgumentError (see above for traceback): Cannot assign a device for operation graph_convolution_feature_steered_convolution_weights/graph_convolution_feature_steered_convolution/utils_convert_to_block_diag_2d/SparseReshape: Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
Registered kernels:
  device='CPU'

I have the same error on windows and linux with tf 1.15 and cuda 10.

[IDSPAGET]

Hi, Do we have an update on this issue? Anything I can do to help tracking the bug?

[avneesh-sud]

Hi @IDSPAGET - can you share a colab or snippet with context on how you are running the graph conv op? Are you running in TFv1 or TFv2? graph or eager?

[avneesh-sud]

Hi @IDSPAGET - confirming the reshape op is executing on CPU, below is the code snippet

import tensorflow as tf
from tensorflow_graphics.geometry.convolution import utils as gc_utils

tf.config.set_soft_device_placement(True)
tf.debugging.set_log_device_placement(True)

with tf.device('/device:GPU:0'):
  st = tf.sparse.SparseTensor(indices=[[0, 0, 0, 1], [1, 0, 0, 2]], values=[1, 2], dense_shape=[2, 2, 3, 4])
  reshaped = tf.sparse.reshape(st, shape=(-1, 3, 4), name=None)

Will engage with TF team, or investigate a GPU friendly alternative.

[IDSPAGET]

Hi @avneesh-sud ,

Thank you for your answer. Here is the example from #13, slightly modified to reproduce the issue.

import numpy as np
import tensorflow as tf
from tensorflow_graphics.nn.layer import graph_convolution as graph_conv

#Number of meshes
N = 2
#Number of spatial dimensions
d = 2
#device to execute with
device='/GPU:0'

#Data consists of the vertices of two meshes.  
# The first mesh has 5 vertices and the second has 4.
# #Shape of data is (numberOfMeshes,maxNumberofVertices,numberofSpatialDimensions)

#An array containing the actual size of each non-padded mesh
sz = np.array([5,4],dtype=np.int64)
#The maximum number of vertices in a mesh
datav = 5

with tf.device(device):
    #Input placeholder for input data (vertices)
    V0 = tf.placeholder(dtype=tf.float64,name="V0",shape=(N,datav,d)) 
    #Input Placeholder for labels for classification (For now, I'm just using throw-away data as my labels)
    L = tf.placeholder(shape=(N,5,1),dtype=tf.float64)
    SZ = tf.placeholder(shape=(N),dtype=tf.int64)
    #Input placeholder for the sparse array representing the adjacency matrix shape:(numberOfMeshes,datav,datav)
    adj_sp = tf.sparse_placeholder(shape=(SZ.shape[0],datav,datav),dtype=tf.float64,name='SPP')
    #The steered graph convolution that is included in Tensorflow's new graphics package
    output = graph_conv.feature_steered_convolution_layer(data=V0,neighbors=adj_sp,sizes=SZ,translation_invariant=False,num_weight_matrices=1,num_output_channels=1)
    #Loss def
    loss = tf.losses.softmax_cross_entropy(L,output, weights=1.0)
    optimisation_step = tf.train.AdamOptimizer(learning_rate=.001).minimize(loss) #Warning not raised if this is commented out
    init_op = tf.initialize_all_variables()

with tf.Session() as sess:
    sess.run(init_op)

This runs fine on cpu with device= '/CPU:0' or with automatic device placement, but yields this error when forcing the device placement on GPU.

InvalidArgumentError: Cannot assign a device for operation graph_convolution_feature_steered_convolution_weights_2/graph_convolution_feature_steered_convolution/utils_convert_to_block_diag_2d/SparseReshape: 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] possible_devices_=[]
SparseReshape: CPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  graph_convolution_feature_steered_convolution_weights_2/graph_convolution_feature_steered_convolution/utils_convert_to_block_diag_2d/SparseReshape (SparseReshape) /device:GPU:0

Op: SparseReshape
Node attrs: 
Registered kernels:
  device='CPU'

     [[{{node graph_convolution_feature_steered_convolution_weights_2/graph_convolution_feature_steered_convolution/utils_convert_to_block_diag_2d/SparseReshape}}]] ```

[IDSPAGET]

@avneesh-sud as pointed out in #13 , the sparse tensors are anyway turned into dense by the optimizer. Maybe a first fix would be to add an option to switch to dense tensors for the graph convolutions instead.