提交 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
...@@ -77,13 +77,20 @@ class CPUDropoutKernel : public framework::OpKernel<T> { ...@@ -77,13 +77,20 @@ class CPUDropoutKernel : public framework::OpKernel<T> {
} }
} }
} else { } 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) { if (upscale_in_train) {
Y.device(place) = X; 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 { } else {
auto X = EigenMatrix<T>::Reshape(*x, 1);
auto Y = EigenMatrix<T>::Reshape(*y, 1);
auto& place =
*context.template device_context<DeviceContext>().eigen_device();
Y.device(place) = X * static_cast<T>(1.0f - dropout_prob); 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.
先完成此消息的编辑!
想要评论请 注册