[EMBED] improve speed

liuyusg · Apr 1, 2016 · 66bd728 · 66bd728
1 parent e75660f
commit 66bd728
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/mshadow b/mshadow
diff --git a/src/operator/embedding-inl.h b/src/operator/embedding-inl.h
@@ -86,7 +86,14 @@ class EmbeddingOp : public Operator {
     Tensor<xpu, 1> data = in_data[embedding::kData].get<xpu, 1, real_t>(s);
     Tensor<xpu, 2> grad_out = out_grad[embedding::kOut].get<xpu, 2, real_t>(s);
     Tensor<xpu, 2> grad_in = in_grad[embedding::kWeight].get<xpu, 2, real_t>(s);
-    Assign(grad_in, req[embedding::kWeight], take_grad(data, grad_out, param_.input_dim));
+    if (req[embedding::kWeight] == kWriteTo) {
+      grad_in = 0.0f;
+      AddTakeGrad(grad_in, data, grad_out);
+    } else if (req[embedding::kWeight] == kAddTo) {
+      AddTakeGrad(grad_in, data, grad_out);
+    } else {
+      LOG(FATAL) << "wrong req";
+    }
   }
 
  private:
+38 −0		mshadow/cuda/tensor_gpu-inl.cuh
+7 −1		mshadow/half.h
+21 −0		mshadow/tensor.h
+9 −0		mshadow/tensor_cpu-inl.h
+6 −0		mshadow/tensor_gpu-inl.h