@@ -49,7 +49,7 @@ void PrepareSamples(const framework::ExecutionContext& context) {
49
49
50
50
int num_label = label_dims.size () == 2 ? label_dims[1 ] : 1 ;
51
51
int index = 0 ;
52
- for (size_t i = 0 ; i < label_dims[0 ]; ++i) {
52
+ for (int64_t i = 0 ; i < label_dims[0 ]; ++i) {
53
53
int j = 0 ;
54
54
for (; j < num_label; ++j) {
55
55
sample_labels_data[index++] = label_data[i * num_label + j];
@@ -86,7 +86,7 @@ class NCEKernel : public framework::OpKernel<T> {
86
86
T* out_data = out->mutable_data <T>(context.GetPlace ());
87
87
int num_neg_samples = context.Attr <int >(" num_neg_samples" );
88
88
int num_total_classes = context.Attr <int >(" num_total_classes" );
89
- int num_true_class = 1 ;
89
+ int64_t num_true_class = 1 ;
90
90
if (label != nullptr ) {
91
91
num_true_class = label->dims ()[1 ];
92
92
}
@@ -95,18 +95,18 @@ class NCEKernel : public framework::OpKernel<T> {
95
95
auto bias = context.Input <Tensor>(" Bias" );
96
96
if (bias != nullptr ) {
97
97
const T* bias_data = bias->data <T>();
98
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
98
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
99
99
sample_out_data[i] = bias_data[sample_labels_data[i]];
100
100
}
101
101
} else {
102
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
102
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
103
103
sample_out_data[i] = 0 ;
104
104
}
105
105
}
106
106
// forward mul
107
107
auto input_mat = EigenMatrix<T>::From (*(context.Input <Tensor>(" Input" )));
108
108
auto weight_mat = EigenMatrix<T>::From (*(context.Input <Tensor>(" Weight" )));
109
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
109
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
110
110
Eigen::Tensor<T, 0 , Eigen::RowMajor, Eigen::DenseIndex> result =
111
111
(input_mat.chip ((int )(i / sample_labels->dims ()[1 ]), 0 ) *
112
112
weight_mat.chip (sample_labels_data[i], 0 ))
@@ -115,8 +115,8 @@ class NCEKernel : public framework::OpKernel<T> {
115
115
sample_out_data[i] = (1 . / (1 . + exp (-sample_out_data[i])));
116
116
}
117
117
// forward cost
118
- for (size_t i = 0 ; i < sample_labels->dims ()[0 ]; ++i) {
119
- size_t j = 0 ;
118
+ for (int64_t i = 0 ; i < sample_labels->dims ()[0 ]; ++i) {
119
+ int64_t j = 0 ;
120
120
out_data[i] = 0 ;
121
121
T w = sample_weight == nullptr ? 1 . : sample_weight_data[i];
122
122
// for true classes
@@ -162,7 +162,7 @@ class NCEGradKernel : public framework::OpKernel<T> {
162
162
T* sample_grad_data =
163
163
sample_grad.mutable_data <T>(sample_labels->dims (), context.GetPlace ());
164
164
// backward cost
165
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
165
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
166
166
T o = sample_out_data[i];
167
167
T w = sample_weight == nullptr
168
168
? 1
@@ -177,7 +177,7 @@ class NCEGradKernel : public framework::OpKernel<T> {
177
177
if (d_bias != nullptr ) {
178
178
T* d_bias_data = d_bias->mutable_data <T>(context.GetPlace ());
179
179
std::fill (d_bias_data, d_bias_data + d_bias->numel (), 0.0 );
180
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
180
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
181
181
d_bias_data[sample_labels_data[i]] += sample_grad_data[i];
182
182
}
183
183
}
@@ -188,7 +188,7 @@ class NCEGradKernel : public framework::OpKernel<T> {
188
188
std::fill (d_w_data, d_w_data + d_w->numel (), 0.0 );
189
189
auto d_w_matrix = EigenMatrix<T>::From (*d_w);
190
190
auto x_matrix = EigenMatrix<T>::From (*(context.Input <Tensor>(" Input" )));
191
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
191
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
192
192
d_w_matrix.chip (sample_labels_data[i], 0 ) +=
193
193
x_matrix.chip ((int )(i / sample_labels->dims ()[1 ]), 0 ) *
194
194
sample_grad_data[i];
@@ -200,7 +200,7 @@ class NCEGradKernel : public framework::OpKernel<T> {
200
200
d_x->mutable_data <T>(context.GetPlace ());
201
201
auto d_x_matrix = EigenMatrix<T>::From (*d_x);
202
202
auto w_matrix = EigenMatrix<T>::From (*(context.Input <Tensor>(" Weight" )));
203
- for (size_t i = 0 ; i < sample_labels->numel (); ++i) {
203
+ for (int64_t i = 0 ; i < sample_labels->numel (); ++i) {
204
204
d_x_matrix.chip ((int )(i / sample_labels->dims ()[1 ]), 0 ) +=
205
205
w_matrix.chip (sample_labels_data[i], 0 ) * sample_grad_data[i];
206
206
}
0 commit comments