提交 e8405e5c 编写于 作者: G GaoWei8 提交者: tensor-tang

Modify the dropout op to multi-thread (#19504)

* Modify the dropout op to multi-thread
test=develop

* define parallel
test=develop
上级 2916caa2
......@@ -76,14 +76,21 @@ class CPUDropoutKernel : public framework::OpKernel<T> {
}
}
}
} else {
if (upscale_in_train) {
const auto* X_data = x->data<T>();
auto* Y_data = y->mutable_data<T>(context.GetPlace());
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int i = 0; i < x->numel(); i++) {
Y_data[i] = X_data[i];
}
} else {
auto X = EigenMatrix<T>::Reshape(*x, 1);
auto Y = EigenMatrix<T>::Reshape(*y, 1);
auto& place =
*context.template device_context<DeviceContext>().eigen_device();
if (upscale_in_train) {
Y.device(place) = X;
} else {
Y.device(place) = X * static_cast<T>(1.0f - dropout_prob);
}
}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册