Skip to content

Do not register Dropout(12) as training ONLY kernel. #3859

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
M. Zeeshan Siddiqui (codemzs) merged 6 commits into from
May 10, 2020
Merged

Do not register Dropout(12) as training ONLY kernel. #3859

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
032a446
Do not register Dropout(12) as training ONLY kernel.
codemzs May 7, 2020
04f5f81
Move Dropout forward implementation in inference project.
codemzs May 8, 2020
fa56141
Merge branch 'master' of https://github.com/microsoft/onnxruntime int…
codemzs May 8, 2020
f69c8fb
fix inference build test failures.
codemzs May 9, 2020
7fde7d1
remove fp16 test since its support is absent on CPU.
codemzs May 9, 2020
dc9683b
build break.
codemzs May 9, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,15 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, uint8_t, ReduceMin);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, GatherND);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Einsum);
// REVIEW(codemzs): ConstEigenVectorArrayMap.cast<MLFLoat16) does not seem to be supported.
// However these types work on GPU implementation.
//class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, MLFloat16_MLFloat16, Dropout);
//class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, MLFloat16_float, Dropout);
//class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, MLFloat16_double, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, float_float, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, float_double, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, double_float, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, double_double, Dropout);

Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
static const BuildKernelCreateInfoFn function_table[] = {
Expand Down Expand Up @@ -1085,6 +1094,15 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
ReduceMin)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, GatherND)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Einsum)>,
// REVIEW(codemzs): ConstEigenVectorArrayMap.cast<MLFLoat16) does not seem to be supported.
// However these types work on GPU implementation.
//BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, MLFloat16_MLFloat16, Dropout)>,
//BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, MLFloat16_float, Dropout)>,
//BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, MLFloat16_double, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, float_float, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, float_double, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, double_float, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, double_double, Dropout)>,
};

for (auto& function_table_entry : function_table) {
Expand Down
32 changes: 32 additions & 0 deletions onnxruntime/core/providers/cpu/nn/dropout_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/providers/cpu/nn/dropout_op.h"

namespace onnxruntime {

// Dropout
#define REGISTER_KERNEL_TYPED(OpName, VER, T1, T2, Trainable) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
OpName, \
kOnnxDomain, \
VER, \
T1##_##T2, \
kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>()), \
Dropout<T1, T2, Trainable>);

// REVIEW(mzs): ConstEigenVectorArrayMap.cast<MLFLoat16) does not seem to be supported.
// However these types work on GPU implementation.
// REGISTER_KERNEL_TYPED(MLFloat16, MLFloat16)
// REGISTER_KERNEL_TYPED(MLFloat16, float)
// REGISTER_KERNEL_TYPED(MLFloat16, double)

REGISTER_KERNEL_TYPED(Dropout, 12, float, float, false)
REGISTER_KERNEL_TYPED(Dropout, 12, float, double, false)
REGISTER_KERNEL_TYPED(Dropout, 12, double, float, false)
REGISTER_KERNEL_TYPED(Dropout, 12, double, double, false)
} // namespace onnxruntime
101 changes: 101 additions & 0 deletions onnxruntime/core/providers/cpu/nn/dropout_op.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "core/framework/op_kernel.h"
#include "core/framework/random_generator.h"
#include <chrono>
#include <random>
#include "core/util/math_cpuonly.h"

namespace onnxruntime {

template <typename T1, typename T2, bool trainable_dropout>
class Dropout final: public OpKernel {
public:
Dropout(const OpKernelInfo& info) : OpKernel{info} {
int64_t seed = 0;
if (info.GetAttr<int64_t>("seed", &seed).IsOK()) {
generator_ = onnxruntime::make_unique<RandomGenerator>(seed);
}
}

Status Compute(OpKernelContext* context) const override;

private:
mutable std::unique_ptr<RandomGenerator> generator_;
};

namespace {
constexpr float k_default_ratio{0.5f};

template <typename T2>
float GetRatioOrDefault(const Tensor* ratio_tensor) {
if (ratio_tensor) {
ORT_ENFORCE(ratio_tensor->Shape().Size() == 1, "ratio input should have a single value.");
#ifdef _WIN32
#pragma warning(disable : 4244)
#endif
const float ratio_value = *ratio_tensor->Data<T2>();
ORT_ENFORCE(0.0f <= ratio_value && ratio_value < 1.0f, "ratio must be in the range [0, 1)");
return ratio_value;
}
return k_default_ratio;
}
} // namespace

template <typename T1, typename T2, bool trainable_dropout>
Status Dropout<T1, T2, trainable_dropout>::Compute(OpKernelContext* context) const {
const Tensor* X = context->Input<Tensor>(0);
auto X_span = X->DataAsSpan<T1>();
const Tensor* ratio = context->Input<Tensor>(1); // optional
const float ratio_value = GetRatioOrDefault<T2>(ratio);
const auto& X_shape = X->Shape();
Tensor* Y = context->Output(0, X_shape);
auto Y_span = Y->MutableDataAsSpan<T1>();
Tensor* mask = context->Output(1, X_shape); // optional
std::unique_ptr<bool[]> temp_mask_buffer{}; // temporary buffer to use if mask input is not provided
auto mask_span = [&X_shape, mask, &temp_mask_buffer]() {
if (mask) return mask->MutableDataAsSpan<bool>();
temp_mask_buffer = onnxruntime::make_unique<bool[]>(X_shape.Size());
return gsl::make_span(temp_mask_buffer.get(), X_shape.Size());
}();

ORT_ENFORCE(!mask || mask->Shape() == X_shape, "X and mask should have the same shape");

const Tensor* training_mode = context->Input<Tensor>(2);
if ((0 == ratio_value /*Backward compat with TrainableDropout*/) ||
!trainable_dropout && (training_mode == nullptr || *(training_mode->Data<bool>()) == false)) {
// drop none
if (X_span.data() != Y_span.data()) {
std::copy(X_span.begin(), X_span.end(), Y_span.begin());
}

if (mask != nullptr) {
std::fill(mask_span.begin(), mask_span.end(), true);
}

} else {
// drop some
ConstEigenVectorArrayMap<T1> X_arr(X_span.data(), X_span.size());
EigenVectorArrayMap<T1> Y_arr(Y_span.data(), Y_span.size());
EigenVectorArrayMap<bool> mask_arr(mask_span.data(), mask_span.size());

// generate mask
{
RandomGenerator& generator = generator_ != nullptr ? *generator_.get() : RandomGenerator::Default();
std::default_random_engine rng(generator.NextSeed());
std::uniform_real_distribution<float> dist{0.0f, 1.0f};
mask_arr = Eigen::ArrayX<bool>::NullaryExpr(
mask_arr.size(),
[ratio_value, &dist, &rng]() { return dist(rng) >= ratio_value; });
}

Y_arr = mask_arr.cast<T1>() * X_arr / (1.0f - ratio_value);
}

return Status::OK();
}

} // namespace onnxruntime
20 changes: 20 additions & 0 deletions onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -773,6 +773,16 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,

class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, int64_t, GatherND);

class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_MLFloat16, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_float, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_double, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_MLFloat16, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_float, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_double, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_MLFloat16, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_float, Dropout);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_double, Dropout);

static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
static const BuildKernelCreateInfoFn function_table[] = {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MemcpyFromHost)>,
Expand Down Expand Up @@ -1289,6 +1299,16 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, uint8_t, ReduceMin)>,

BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, int64_t, GatherND)>,

BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_MLFloat16, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_float, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_double, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_MLFloat16, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_float, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_double, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_MLFloat16, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_float, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_double, Dropout)>,
};

for (auto& function_table_entry : function_table) {
Expand Down
35 changes: 35 additions & 0 deletions onnxruntime/core/providers/cuda/nn/dropout.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/providers/cuda/nn/dropout.h"

namespace onnxruntime {
namespace cuda {

#define REGISTER_KERNEL_TYPED(T1, T2) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
Dropout, \
kOnnxDomain, \
12, \
T1##_##T2, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>()) \
.InputMemoryType<OrtMemTypeCPUInput>(1) \
.InputMemoryType<OrtMemTypeCPUInput>(2), \
Dropout<T1, T2, false>);

REGISTER_KERNEL_TYPED(MLFloat16, MLFloat16)
REGISTER_KERNEL_TYPED(MLFloat16, float)
REGISTER_KERNEL_TYPED(MLFloat16, double)
REGISTER_KERNEL_TYPED(float, MLFloat16)
REGISTER_KERNEL_TYPED(float, float)
REGISTER_KERNEL_TYPED(float, double)
REGISTER_KERNEL_TYPED(double, MLFloat16)
REGISTER_KERNEL_TYPED(double, float)
REGISTER_KERNEL_TYPED(double, double)

} // namespace cuda
} // namespace onnxruntime
95 changes: 95 additions & 0 deletions onnxruntime/core/providers/cuda/nn/dropout.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/nn/dropout_impl.h"
#include "core/providers/cuda/nn/dropout.h"
#include "core/providers/common.h"
#include "core/framework/random_seed.h"

namespace onnxruntime {
namespace cuda {

template <typename T1, typename T2, bool trainable_dropout>
class Dropout final : public CudaKernel {
public:
Dropout(const OpKernelInfo& info) : CudaKernel(info), default_ratio_(0.5) {
int64_t seed = 0;
if (info.GetAttr<int64_t>("seed", &seed).IsOK()) {
generator_ = onnxruntime::make_unique<PhiloxGenerator>(static_cast<uint64_t>(seed));
}
}

Status ComputeInternal(OpKernelContext* context) const override;

private:
mutable std::unique_ptr<PhiloxGenerator> generator_;
const float default_ratio_;
};

template <typename T1, typename T2, bool trainable_dropout>
Status Dropout<T1, T2, trainable_dropout>::ComputeInternal(OpKernelContext* context) const {
typedef typename ToCudaType<T1>::MappedType CudaT;

//Get X_data
const Tensor* X = context->Input<Tensor>(0);
if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
const TensorShape& shape = X->Shape();
auto X_data = reinterpret_cast<const CudaT*>(X->template Data<T1>());
const int64_t N = shape.Size();

//Get Y_data
auto Y = context->Output(0, shape);
auto Y_data = reinterpret_cast<CudaT*>(Y->template MutableData<T1>());

//Get mask_data
auto mask = context->Output(1, shape);
ORT_ENFORCE(!mask || mask->Shape().Size() == N);

//Get the ratio_data
float ratio_data;
auto ratio = context->Input<Tensor>(1);

static_assert(std::is_same<T2, MLFloat16>::value || std::is_same<T2, float>::value || std::is_same<T2, double>::value,
"T2 must be float16 or float or double");

if (ratio) {
ratio_data = static_cast<float>(*(ratio->template Data<T2>()));
} else {
ratio_data = default_ratio_;
}
ORT_ENFORCE(ratio_data >= 0.0f && ratio_data < 1.0f);

const Tensor* training_mode = context->Input<Tensor>(2);
//Check for inference mode.
if ((0 == ratio_data /*Backward compat with TrainableDropout*/) ||
(!trainable_dropout && (training_mode == nullptr || *(training_mode->Data<bool>()) == false))) {
if (Y_data != X_data) {
CUDA_CALL_THROW(cudaMemcpyAsync(Y_data, X_data, N * sizeof(T1), cudaMemcpyDeviceToDevice));
}

// If mask is requested, return all 1s.
if (mask != nullptr) {
ORT_ENFORCE(cudaMemset(mask->MutableData<bool>(), true, N * sizeof(bool)) == cudaSuccess);
}

return Status::OK();
}

IAllocatorUniquePtr<bool> temp_mask_buffer{}; // buffer to use if mask is not provided
bool* const mask_data = [this, N, mask, &temp_mask_buffer]() {
if (mask) return mask->MutableData<bool>();
temp_mask_buffer = GetScratchBuffer<bool>(N);
return temp_mask_buffer.get();
}();

PhiloxGenerator& generator = generator_ != nullptr ? *generator_.get() : PhiloxGenerator::Default();
DropoutKernelImpl(GetDeviceProp(), N, ratio_data, generator, X_data, Y_data, mask_data);

return Status::OK();
}

} // namespace cuda
} // namespace onnxruntime
Loading