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

Add SoftmaxCrossEntropyLoss to mixed-precision-transformer. (#3760)

Browse source
This commit is contained in:
M. Zeeshan Siddiqui 2020-04-30 02:48:21 -07:00 committed by GitHub
parent 9f72752397
commit b9a5ed1fe2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
orttraining/orttraining/core/graph/mixed_precision_transformer.cc
View file

@ -33,7 +33,9 @@ namespace training {
// continue to use 32-bit precision. Others will used reduced precision.
static const std::unordered_set<std::string> FP32_Nodes = {
"SparseSoftmaxCrossEntropy",
"SparseSoftmaxCrossEntropyGrad"};
"SparseSoftmaxCrossEntropyGrad",
"SoftmaxCrossEntropyLoss",
"SoftmaxCrossEntropyLossGrad"};
bool IsFP32Node(const Node* node) {
return FP32_Nodes.find(node->OpType()) != FP32_Nodes.cend();
@ -54,6 +56,8 @@ static const std::unordered_map<std::string, std::vector<int>> stage2_fp32_node_
{"DropoutGrad", {2}},
{"SparseSoftmaxCrossEntropy", {0, 2}},
{"SparseSoftmaxCrossEntropyGrad", {0, 1, 3}},
{"SoftmaxCrossEntropyLoss", {0, 2}},
{"SoftmaxCrossEntropyLossGrad", {0, 1, 3}},
};
bool IsFP32(const std::unordered_map<std::string, std::vector<int>>& map, std::string opname, int argnum) {