Repository: Theano/Theano

Issue:
CrossentropyCategorical1Hot() not on GPU
I have see a case where CrossentropyCategorical1Hot was not on the GPU. We don't have it implemented, just CrossentropyCategorical1Hot with the softmax was implemented on the GPU.

It happen that there is something between the softmax and the crossentropy. This currently make this not working on the GPU.

As we implemented logsoftmax on the CPU and GPU, I think we can just implement CrossentropyCategorical1Hot on the GPU and we will support those cases while staying numericaly stable?



Discussion:
That Op is not really needed, as it does not provide numerical stability. In fact, it could actually prevent LogSoftmax from being inserted.
It is only useful as an intermediate step before the bigger CrossEntropySoftmax... gets inserted, and maybe for the special case of advanced indexing it is performing.
I think we should replace it by `-log(coding_dist[arange(true_one_of_n.shape[0]), true_one_of_n])` or something like that.
We would need to wait for advanced indexing on the GPU for that, though

It is actually still useful when `coding_dist` is not the output of a softmax.
I have a PR in preparation for that.

Resolving patch:
diff --git a/theano/gpuarray/opt.py b/theano/gpuarray/opt.py
--- a/theano/gpuarray/opt.py
+++ b/theano/gpuarray/opt.py
@@ -19,6 +19,7 @@
 from theano.misc.ordered_set import OrderedSet
 
 from theano.scalar.basic import Scalar, Pow, Cast
+from theano.scalar.basic import log, neg, true_div
 from theano.scalar.basic_scipy import Erfinv, Erfcinv
 from theano.scan_module import scan_utils, scan_op, scan_opt
 
@@ -163,6 +164,8 @@ def f(local_opt):
                        'unsafe')
 
 
+# Define a few operations to use in optimizations,
+# in order to avoid introducin new CPU Ops, or useless ones.
 def safe_to_gpu(x, ctx_name):
     if isinstance(x.type, tensor.TensorType):
         return GpuFromHost(ctx_name)(x)
@@ -176,6 +179,10 @@ def safe_to_cpu(x):
     else:
         return x
 
+gpu_log = GpuElemwise(log)
+gpu_neg = GpuElemwise(neg)
+gpu_true_div = GpuElemwise(true_div)
+
 
 def op_lifter(OP, cuda_only=False):
     """
@@ -1321,6 +1328,38 @@ def local_gpua_softmaxwithbias(op, context_name, inputs, outputs):
     return gpu_softmax_with_bias
 
 
+@register_opt('fast_compile')
+@op_lifter([tensor.nnet.CrossentropyCategorical1Hot])
+@register_opt2([tensor.nnet.CrossentropyCategorical1Hot], 'fast_compile')
+def local_gpu_crossentropycategorical1hot(op, context_name, inputs, outputs):
+    # There is no corresponding GPU Op, but we can express it as:
+    #   coding, one_of_n = inputs
+    #   -log(coding[arange(coding.shape[0]), one_of_n])
+    coding, one_of_n = inputs
+    idx0 = theano.tensor.arange(shape_i(coding, 0))
+    return [gpu_neg(gpu_log(coding[idx0, one_of_n]))]
+
+
+@register_opt('fast_compile')
+@op_lifter([tensor.nnet.CrossentropyCategorical1HotGrad])
+@register_opt2([tensor.nnet.CrossentropyCategorical1HotGrad], 'fast_compile')
+def local_gpu_crossentropycategorical1hotgrad(op, context_name, inputs, outputs):
+    # There is no corresponding GPU Op, but we can express it as:
+    #   gy, coding, one_of_n = inputs
+    #   gcoding = zeros_like(coding)
+    #   gcoding[arange(coding.shape[0]), one_of_n] = -g / (
+    #       coding[arange(coding.shape[0]), one_of_n])
+    gy, coding, one_of_n = inputs
+    idx0 = theano.tensor.arange(shape_i(coding, 0))
+    z = GpuAlloc(context_name, memset_0=True)(
+        as_gpuarray_variable(np.zeros((), dtype=coding.dtype), context_name),
+        *[shape_i(coding, i) for i in xrange(coding.ndim)])
+    gcoding = tensor.set_subtensor(
+        z[idx0, one_of_n],
+        gpu_neg(gpu_true_div(gy, coding[idx0, one_of_n])))
+    return [gcoding.transfer(context_name)]
+
+
 @register_opt('fast_compile')
 @op_lifter([theano.tensor.opt.Assert])
 def local_gpua_assert(op, context_name, inputs, outputs):


Test patch:
diff --git a/theano/gpuarray/tests/test_opt.py b/theano/gpuarray/tests/test_opt.py
--- a/theano/gpuarray/tests/test_opt.py
+++ b/theano/gpuarray/tests/test_opt.py
@@ -680,3 +680,17 @@ def randX(*args):
         z = tensor.batched_dot(x, y)
         f = theano.function([x, y], z, mode=mode_with_gpu)
         f(x_val, y_val)
+
+
+def test_crossentropycategorical1hot_lifter():
+    rng = np.random.RandomState(utt.fetch_seed())
+    x = tensor.matrix()
+    y = tensor.lvector()
+    z = tensor.nnet.crossentropy_categorical_1hot(x, y)
+    gx = theano.grad(z.mean(), x)
+    f = theano.function([x, y], [z, gx], mode=mode_with_gpu)
+    assert not any(isinstance(n.op, (tensor.nnet.CrossentropyCategorical1Hot,
+                                     tensor.nnet.CrossentropyCategorical1HotGrad))
+                   for n in f.maker.fgraph.apply_nodes)
+    f(rng.uniform(0.1, 0.9, (13, 5)).astype(theano.config.floatX),
+      rng.randint(5, size=(13,)))