efficient implementation of softlabel softmax loss layer (GPU)

src/cuda/backend.jl

@@ -50,6 +50,8 @@ end
   logistic_loss_forward_double,
   softmax_loss_backward_float,
   softmax_loss_backward_double,
+  softlabel_softmax_loss_backward_float,
+  softlabel_softmax_loss_backward_double,
   relu_forward_float,
   relu_forward_double,
   relu_backward_float,

src/cuda/kernels/kernels.cu

@@ -10,6 +10,7 @@
 
 #include "logistic_loss.impl"
 #include "softmax_loss.impl"
+#include "softlabel_softmax_loss.impl"
 #include "accuracy.impl"
 #include "channel_pooling.impl"
 #include "dropout.impl"

src/cuda/kernels/softlabel_softmax_loss.impl

@@ -0,0 +1,22 @@
+template <typename T>
+__device__ void softlabel_softmax_loss_backward(T *prob, T *label, int num, int spatial_dim, int prob_dim) {
+  int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  int idx_sp = blockIdx.y;
+  int idx_y = blockIdx.z;
+  if (idx >= num || idx_sp >= spatial_dim || idx_y >= prob_dim)
+    return;
+
+  int loc = idx * (spatial_dim*prob_dim) + idx_y*spatial_dim + idx_sp;
+  prob[loc] -= label[loc];
+}
+
+extern "C" {
+  __global__ void softlabel_softmax_loss_backward_float(float *prob, float *label, int num, int spatial_dim, int prob_dim) {
+    softlabel_softmax_loss_backward(prob, label, num, spatial_dim, prob_dim);
+  }
+  __global__ void softlabel_softmax_loss_backward_double(double *prob, double *label, int num, int spatial_dim, int prob_dim) {
+    softlabel_softmax_loss_backward(prob, label, num, spatial_dim, prob_dim);
+  }
+}
+
+// vim: ft=cuda

src/cuda/layers/softlabel-softmax-loss.jl

@@ -2,16 +2,27 @@ function backward(backend::GPUBackend, state::SoftlabelSoftmaxLossLayerState, in
   diff = diffs[1]
 
   if isa(diff, CuTensorBlob)
-    pred = inputs[1]
-    label = inputs[2]
-    data_type = eltype(pred)
-    copy!(state.softmax_loss.logistic.weights_blob, label)
-    erase!(diff)
+    copy!(diff, state.softmax_loss.softmax.blobs[1])
 
-    for i = 1:length(state.fake_labels)
-      backward(backend, state.softmax_loss, Blob[pred, state.fake_labels[i]], Blob[state.fake_diff])
-      CuBLAS.axpy(backend.cublas_ctx, length(pred), one(data_type), state.fake_diff.ptr, 1, diff.ptr, 1)
+    data_type = eltype(diff)
+    spatial_dim, channels, num = split_dims(diff, state.softmax_loss.logistic.op_dim)
+    prob_dim = channels
+
+    x_block = int(ceil(float64(num)/CUDA.THREADS_PER_BLOCK_X))
+    y_block = spatial_dim
+    z_block = prob_dim
+
+    if data_type == Float32
+      kernel = backend.mocha.softlabel_softmax_loss_backward_float
+    elseif data_type == Float64
+      kernel = backend.mocha.softlabel_softmax_loss_backward_double
+    else
+      error("Unsupported data type $data_type")
     end
+    CUDA.launch(kernel, (x_block, y_block, z_block), (CUDA.THREADS_PER_BLOCK_X, 1, 1),
+        (diff.ptr.p, inputs[2].ptr.p, num, spatial_dim, prob_dim))
+    CuBLAS.scal(backend.cublas_ctx, length(diff), convert(data_type, 1.0/(spatial_dim*num)),
+        diff.ptr, 1)
   end
 end