saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Prevent divide by zero in CUDA implementation of SoftmaxCrossEntropyLossGrad. (#3962)

Browse source
This commit is contained in:
M. Zeeshan Siddiqui 2020-05-16 00:33:25 -07:00 committed by GitHub
parent 132ce3a561
commit a296b16719
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu
View file

@ -117,7 +117,11 @@ __global__ void _WeightedSoftmaxCrossEntropyLossGrad(
int row = i / C;
int d = i % C;
CUDA_KERNEL_ASSERT(weight[row] == 0 || (label[row] >= 0 && label[row] < C));
output_data[i] = (*dY) * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor);
if(0 == *normalize_factor){
output_data[i] = 0;
} else {
output_data[i] = (*dY) * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor);
}
}
template <typename T, typename Tin>
@ -135,7 +139,11 @@ __global__ void _WeightedReductionNoneSoftmaxCrossEntropyLossGrad(
int row = i / C;
int d = i % C;
CUDA_KERNEL_ASSERT(weight[row] == 0 || (label[row] >= 0 && label[row] < C));
output_data[i] = dY[row] * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor);
if(0 == *normalize_factor){
output_data[i] = 0;
} else {
output_data[i] = dY[row] * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor);
}
}
template <typename T, typename Tin>