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
|
#include "caffe2/operators/sparse_dropout_with_replacement_op.h"
#include <algorithm>
#include <iterator>
namespace caffe2 {
template <>
bool SparseDropoutWithReplacementOp<CPUContext>::RunOnDevice() {
auto& X = Input(0);
CAFFE_ENFORCE_EQ(X.ndim(), 1, "Input tensor should be 1-D");
const int64_t* Xdata = X.data<int64_t>();
auto& Lengths = Input(1);
CAFFE_ENFORCE_EQ(Lengths.ndim(), 1, "Lengths tensor should be 1-D");
auto* OutputLengths = Output(1, Lengths.size(), at::dtype<int32_t>());
int32_t const* input_lengths_data = Lengths.template data<int32_t>();
int32_t* output_lengths_data =
OutputLengths->template mutable_data<int32_t>();
// Check that input lengths add up to the length of input data
int total_input_length = 0;
for (int i = 0; i < Lengths.numel(); ++i) {
total_input_length += input_lengths_data[i];
}
CAFFE_ENFORCE_EQ(
total_input_length,
X.numel(),
"Inconsistent input data. Number of elements should match total length.");
at::bernoulli_distribution<double> dist(1. - ratio_);
auto* gen = context_.RandGenerator();
int32_t total_output_length = 0;
vector<bool> selected(Lengths.numel(), true);
for (int i = 0; i < Lengths.numel(); ++i) {
if (dist(gen) > 0.5) {
output_lengths_data[i] = input_lengths_data[i];
} else {
// Replace with a single dropout value. Even if input length is 0.
output_lengths_data[i] = 1;
selected[i] = false;
}
total_output_length += output_lengths_data[i];
}
auto* Y = Output(0, {total_output_length}, at::dtype<int64_t>());
int64_t* Ydata = Y->template mutable_data<int64_t>();
int input_index = 0;
int output_index = 0;
for (int i = 0; i < Lengths.numel(); ++i) {
if (selected[i]) {
// Copy logical elements from input to output
for (int j = input_index; j < input_index + input_lengths_data[i]; ++j) {
Ydata[output_index++] = Xdata[j];
}
} else {
Ydata[output_index++] = replacement_value_;
}
input_index += input_lengths_data[i];
}
return true;
}
REGISTER_CPU_OPERATOR(
SparseDropoutWithReplacement,
SparseDropoutWithReplacementOp<CPUContext>);
OPERATOR_SCHEMA(SparseDropoutWithReplacement)
.NumInputs(2)
.SameNumberOfOutput()
.SetDoc(R"DOC(
`SparseDropoutWithReplacement` takes a 1-d input tensor and a lengths tensor.
Values in the Lengths tensor represent how many input elements consitute each
example in a given batch. The set of input values for an example will be
replaced with the single dropout value with probability given by the `ratio`
argument.
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"SparseDropoutWithReplacement",
["X", "Lengths"],
["Y", "OutputLengths"],
ratio=0.5,
replacement_value=-1
)
workspace.FeedBlob("X", np.array([1, 2, 3, 4, 5]).astype(np.int64))
workspace.FeedBlob("Lengths", np.array([2, 3]).astype(np.int32))
print("X:", workspace.FetchBlob("X"))
print("Lengths:", workspace.FetchBlob("Lengths"))
workspace.RunOperatorOnce(op)
print("Y:", workspace.FetchBlob("Y"))
print("OutputLengths:", workspace.FetchBlob("OutputLengths"))
```
**Result**
```
X: [1, 2, 3, 4, 5]
Lengths: [2, 3]
Y: [1, 2, -1]
OutputLengths: [2, 1]
```
</details>
)DOC")
.Arg(
"ratio",
"*(type: float; default: 0.0)* Probability of an element to be replaced.")
.Arg(
"replacement_value",
"*(type: int64_t; default: 0)* Value elements are replaced with.")
.Input(0, "X", "*(type: Tensor`<int64_t>`)* Input data tensor.")
.Input(
1,
"Lengths",
"*(type: Tensor`<int32_t>`)* Lengths tensor for input.")
.Output(0, "Y", "*(type: Tensor`<int64_t>`)* Output tensor.")
.Output(1, "OutputLengths", "*(type: Tensor`<int32_t>`)* Output tensor.");
NO_GRADIENT(SparseDropoutWithReplacement);
} // namespace caffe2
|
