tensorflow: contrib.learn.Estimator does not work with multiple GPU

Attempting to assign ops to a GPU within model_fn passed to an Estimator produces the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device to node 'save/ShardedFilename_2': Could not satisfy explicit device specification '/device:GPU:1' because no supported kernel for GPU devices is available.
Colocation Debug Info:
Colocation group had the following types and devices: 
Identity: CPU 
ShardedFilename: CPU 
	 [[Node: save/ShardedFilename_2 = ShardedFilename[_device="/device:GPU:1"](save/Const, save/ShardedFilename_2/shard, save/num_shards)]]

This can be reproduced by running the example in examples/learn/multiple_gpu.py

About this issue

Original URL
State: closed
Created 8 years ago
Comments: 23 (7 by maintainers)

Most upvoted comments

I don’t think creating sparse string Tensors on the GPU will work since the GPU does not support strings.

It’s best to structure your model so the string to int conversion happens on the CPU, and the GPU just processes the dense part of your model.

That said, the linear regression canned estimator will probably see no benefit from running on the GPU (not enough matrix multiplications to offset the data transfer cost).

alextp on Apr 5, 2017

I believe if you did that, you should be fine.

martinwicke on Dec 8, 2016