forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Dropout.cpp
175 lines (147 loc) · 5.83 KB
/
Dropout.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include <ATen/core/Tensor.h>
#include <ATen/NamedTensorUtils.h>
#include <ATen/TensorOperators.h>
#include <c10/util/irange.h>
#ifndef AT_PER_OPERATOR_HEADERS
#include <ATen/Functions.h>
#include <ATen/NativeFunctions.h>
#else
#include <ATen/ops/alpha_dropout_native.h>
#include <ATen/ops/dropout_native.h>
#include <ATen/ops/empty_like.h>
#include <ATen/ops/feature_alpha_dropout_native.h>
#include <ATen/ops/feature_dropout_native.h>
#include <ATen/ops/native_dropout.h>
#include <ATen/ops/native_dropout_backward_native.h>
#include <ATen/ops/native_dropout_native.h>
#include <ATen/ops/ones_like.h>
#include <ATen/ops/zeros.h>
#endif
namespace at::native {
namespace {
template<bool inplace>
using Ctype = typename std::conditional<inplace, Tensor&, Tensor>::type;
Tensor make_feature_noise(const Tensor& input) {
auto input_sizes = input.sym_sizes();
TORCH_CHECK(input.dim() >= 2, "Feature dropout requires at least 2 dimensions in the input");
c10::SymDimVector sizes;
sizes.reserve(input.dim());
sizes.push_back(input_sizes[0]);
sizes.push_back(input_sizes[1]);
for (C10_UNUSED const auto i : c10::irange(2, input.dim())) {
sizes.push_back(1);
}
return input.new_empty_symint(sizes);
}
bool is_fused_kernel_acceptable(const Tensor& input, double p) {
return (input.is_cuda() || input.is_xpu() || input.is_lazy() || input.is_privateuseone()) && p > 0 && p < 1 && input.sym_numel() > 0;
}
// NB: sure, we could have used different overloads here, but I would feel insecure
// knowing that this dispatch depends only on the constness of the references
template<bool inplace>
Tensor& multiply(Tensor& input, const Tensor& noise) {
static_assert(inplace, "Wrong multiply overload triggered in Dropout.cpp");
return input.mul_(noise);
}
template<bool inplace>
Tensor multiply(const Tensor& input, const Tensor& noise) {
static_assert(!inplace, "Wrong multiply overload triggered in Dropout.cpp");
return input.mul(noise);
}
template<bool feature_dropout, bool alpha_dropout, bool inplace, typename T>
Ctype<inplace> _dropout_impl(T& input, double p, bool train) {
TORCH_CHECK(p >= 0 && p <= 1, "dropout probability has to be between 0 and 1, but got ", p);
if (p == 0 || !train || input.sym_numel() == 0) {
return input;
}
if (p == 1) {
return multiply<inplace>(input, at::zeros({}, input.options()));
}
at::Tensor b; // used for alpha_dropout only
auto noise = feature_dropout ? make_feature_noise(input) : at::empty_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
noise.bernoulli_(1 - p);
if (alpha_dropout) {
constexpr double alpha = 1.7580993408473766;
double a = 1. / std::sqrt((alpha * alpha * p + 1) * (1 - p));
b = noise.add(-1).mul_(alpha * a).add_(alpha * a * p);
noise.mul_(a);
} else {
noise.div_(1 - p);
}
if (!alpha_dropout) {
return multiply<inplace>(input, noise);
} else {
return multiply<inplace>(input, noise).add_(b);
}
}
#define ALIAS_SPECIALIZATION(ALIAS_NAME, IS_FEATURE, IS_ALPHA) \
template <bool inplace, typename... Args> \
Ctype<inplace> ALIAS_NAME(Args&&... args) { \
return _dropout_impl<IS_FEATURE, IS_ALPHA, inplace>(std::forward<Args>(args)...); \
}
ALIAS_SPECIALIZATION(_dropout, false, false)
ALIAS_SPECIALIZATION(_feature_dropout, true, false)
ALIAS_SPECIALIZATION(_alpha_dropout, false, true )
ALIAS_SPECIALIZATION(_feature_alpha_dropout, true, true )
} // anomymous namepsace
std::tuple<Tensor,Tensor>
native_dropout_cpu(const Tensor& input, double p, c10::optional<bool> train) {
if (input.numel() == 0) {
return std::make_tuple(input, at::empty_like(input, input.options()));
}
Tensor mask;
Tensor output;
if (!train.has_value() || *train) {
double p1m = 1. - p;
// Check for probability of zero to avoid divide by zero and NaN results
double scale = p1m == 0 ? 0. : 1. / p1m;
mask = at::empty_like(input, input.options().dtype(c10::CppTypeToScalarType<bool>::value));
mask.bernoulli_(p1m);
output = input.mul(mask).mul_(scale);
} else {
mask = at::ones_like(input, input.options().dtype(c10::CppTypeToScalarType<bool>::value));
output = input.clone();
}
return std::make_tuple(output, mask);
}
Tensor native_dropout_backward(const Tensor& grad, const Tensor& mask, double scale) {
Tensor result = grad * mask * scale;
return result;
}
Tensor dropout(const Tensor& input, double p, bool train) {
auto result = [&]() {
NoNamesGuard guard;
// TODO: we can remove this is_nested() code smell in the future
// if we find a way to support _dropout for nested tensor
// e.g. make it an op (at::_dropout) to use dispatcher?
if (input.is_nested() || (train && is_fused_kernel_acceptable(input, p))) {
return std::get<0>(at::native_dropout(input, p, train));
}
return _dropout<false>(input, p, train);
}();
namedinference::propagate_names(result, input);
return result;
}
Tensor& dropout_(Tensor& input, double p, bool train) {
return _dropout<true>(input, p, train);
}
Tensor feature_dropout(const Tensor& input, double p, bool train) {
return _feature_dropout<false>(input, p, train);
}
Tensor& feature_dropout_(Tensor& input, double p, bool train) {
return _feature_dropout<true>(input, p, train);
}
Tensor alpha_dropout(const Tensor& input, double p, bool train) {
return _alpha_dropout<false>(input, p, train);
}
Tensor& alpha_dropout_(Tensor& input, double p, bool train) {
return _alpha_dropout<true>(input, p, train);
}
Tensor feature_alpha_dropout(const Tensor& input, double p, bool train) {
return _feature_alpha_dropout<false>(input, p, train);
}
Tensor& feature_alpha_dropout_(Tensor& input, double p, bool train) {
return _feature_alpha_dropout<true>(input, p, train);
}
} // namespace at::native