dropout_op.cu 3.2 KB
Newer Older
X
Xinghai Sun 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

L
Luo Tao 已提交
3 4 5
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
X
Xinghai Sun 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
X
Xinghai Sun 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
X
Xinghai Sun 已提交
14 15

#define EIGEN_USE_GPU
16 17 18 19
#include <thrust/device_ptr.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/random.h>
#include <thrust/transform.h>
Y
Yi Wang 已提交
20
#include "paddle/fluid/operators/dropout_op.h"
X
Xinghai Sun 已提交
21

22 23 24
namespace paddle {
namespace operators {

25
template <typename T, typename AttrType>
D
dzhwinter 已提交
26 27 28 29 30 31
__global__ void RandomGenerator(const size_t n, const int seed,
                                const AttrType dropout_prob, const T* src,
                                T* mask_data, T* dst) {
  thrust::minstd_rand rng;
  rng.seed(seed);
  thrust::uniform_real_distribution<AttrType> dist(0, 1);
32

D
dzhwinter 已提交
33 34
  int idx = blockDim.x * blockIdx.x + threadIdx.x;
  for (; idx < n; idx += blockDim.x * gridDim.x) {
35
    if (dist(rng) < dropout_prob) {
D
dzhwinter 已提交
36 37 38
      mask_data[idx] = static_cast<T>(0);
    } else {
      mask_data[idx] = static_cast<T>(1);
39
    }
D
dzhwinter 已提交
40
    dst[idx] = mask_data[idx] * src[idx];
41
  }
D
dzhwinter 已提交
42
}
43 44 45 46

// It seems that Eigen::Tensor::setRandom in GPU will SEGFAULT.
// Use std::random and thrust::random(thrust is a std library in CUDA) to
// implement uniform random.
47
template <typename Place, typename T, typename AttrType>
Y
Yu Yang 已提交
48
class GPUDropoutKernel : public framework::OpKernel<T> {
49 50 51 52 53
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    auto* x = context.Input<Tensor>("X");
    auto* y = context.Output<Tensor>("Out");
    y->mutable_data<T>(context.GetPlace());
54
    AttrType dropout_prob = context.Attr<AttrType>("dropout_prob");
55

56 57
    auto X = EigenMatrix<T>::Reshape(*x, 1);
    auto Y = EigenMatrix<T>::Reshape(*y, 1);
58

Q
QI JUN 已提交
59
    auto& place = *context.template device_context<Place>().eigen_device();
60
    if (!context.Attr<bool>("is_test")) {
61 62
      auto* mask = context.Output<Tensor>("Mask");
      auto* mask_data = mask->mutable_data<T>(context.GetPlace());
D
dzhwinter 已提交
63 64 65
      size_t size = framework::product(mask->dims());
      auto* x_data = x->data<T>();
      auto* y_data = y->mutable_data<T>(context.GetPlace());
66 67 68 69 70

      std::random_device rnd;
      int seed =
          context.Attr<bool>("fix_seed") ? context.Attr<int>("seed") : rnd();

D
dzhwinter 已提交
71 72 73 74 75
      int threads = 512;
      int grid = (x->numel() + threads - 1) / threads;
      RandomGenerator<T, AttrType><<<grid, threads, 0,
                                     context.cuda_device_context().stream()>>>(
          size, seed, dropout_prob, x_data, mask_data, y_data);
76
    } else {
77
      Y.device(place) = X * (1.0f - dropout_prob);
78
    }
79 80 81 82 83 84
  }
};

}  // namespace operators
}  // namespace paddle

X
Xinghai Sun 已提交
85
namespace ops = paddle::operators;
Q
QI JUN 已提交
86 87 88 89 90 91
REGISTER_OP_CUDA_KERNEL(
    dropout,
    ops::GPUDropoutKernel<paddle::platform::CUDADeviceContext, float, float>);
REGISTER_OP_CUDA_KERNEL(
    dropout_grad,
    ops::DropoutGradKernel<paddle::platform::CUDADeviceContext, float>);