Add softmax to the mnist example (#4149)

This commit is contained in:
Changming Sun 2020-06-08 09:33:50 -07:00 committed by GitHub
parent 4e1dac67cd
commit 5a5f44eed7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -5,11 +5,25 @@
#include <windowsx.h>
#include <onnxruntime_cxx_api.h>
#include <array>
#include <cmath>
#pragma comment(lib, "user32.lib")
#pragma comment(lib, "gdi32.lib")
#pragma comment(lib, "onnxruntime.lib")
template <typename T>
static void softmax(T& input) {
float rowmax = *std::max_element(input.begin(), input.end());
std::vector<float> y(input.size());
float sum = 0.0f;
for (size_t i = 0; i != input.size(); ++i) {
sum += y[i] = std::exp(input[i] - rowmax);
}
for (size_t i = 0; i != input.size(); ++i) {
input[i] = y[i] / sum;
}
}
// This is the structure to interface with the MNIST model
// After instantiation, set the input_image_ data to be the 28x28 pixel image of the number to recognize
// Then call Run() to fill in the results_ data with the probabilities of each
@ -26,7 +40,7 @@ struct MNIST {
const char* output_names[] = {"Plus214_Output_0"};
session_.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1);
softmax(results_);
result_ = std::distance(results_.begin(), std::max_element(results_.begin(), results_.end()));
return result_;
}