ohem for 2d softmax

dsqx71 · Aug 5, 2016 · 675ffaa · 675ffaa
1 parent 3954709
commit 675ffaa
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/operator/softmax_output-inl.h b/src/operator/softmax_output-inl.h
@@ -196,6 +196,11 @@ class SoftmaxOutputOp : public Operator {
         valid_cnt = 1;
       }
       grad *= DType(param_.grad_scale / valid_cnt);
+      if (param_.is_hidden_layer) {
+        Tensor<xpu, 2, DType> o_grad =
+          out_grad[softmaxout_enum::kOut].get_with_shape<xpu, 2, DType>(s2, s);
+        grad *= o_grad;
+      }
     }
   }