1/**
2 * Copyright (c) Glow Contributors. See CONTRIBUTORS file.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include "Loader.h"
18
19#include "glow/Base/Image.h"
20#include "glow/Converter/TypeAToTypeBFunctionConverter.h"
21#include "glow/Graph/Nodes.h"
22#include "glow/Importer/Caffe2ModelLoader.h"
23#include "glow/Importer/ONNXModelLoader.h"
24#include "glow/Support/Support.h"
25
26#include "llvm/ADT/StringSwitch.h"
27#include "llvm/Support/CommandLine.h"
28#include "llvm/Support/Format.h"
29#include "llvm/Support/Timer.h"
30#include "llvm/Support/raw_ostream.h"
31
32#include <atomic>
33#include <cfloat>
34#include <fstream>
35#include <future>
36#include <iostream>
37#include <memory>
38#include <mutex>
39#include <queue>
40#include <sstream>
41#include <thread>
42
43#include "ExecutorCore.h"
44#include "ExecutorCoreHelperFunctions.h"
45
46using namespace glow;
47
48namespace {
49
50/// Image loader options.
51llvm::cl::OptionCategory imageClassifierCat("Image Loader Options");
52
53llvm::cl::opt<unsigned> labelOffset(
54 "label-offset",
55 llvm::cl::desc("Label offset for TF ONNX models with 1001 classes"),
56 llvm::cl::Optional, llvm::cl::init(0), llvm::cl::cat(imageClassifierCat));
57
58llvm::cl::opt<bool>
59 computeSoftmax("compute-softmax",
60 llvm::cl::desc("Compute softmax of the network output"),
61 llvm::cl::Optional, llvm::cl::init(false),
62 llvm::cl::cat(imageClassifierCat));
63
64llvm::cl::opt<unsigned>
65 topKCount("topk",
66 llvm::cl::desc("Number of highest likelihood labels to print and "
67 "match the correspondent expected-labels"),
68 llvm::cl::Optional, llvm::cl::init(1),
69 llvm::cl::cat(imageClassifierCat));
70
71llvm::cl::list<unsigned> expectedMatchingLabels(
72 "expected-labels",
73 llvm::cl::desc("The comma delimited list of the matching labels"),
74 llvm::cl::value_desc("int"), llvm::cl::ZeroOrMore, llvm::cl::CommaSeparated,
75 llvm::cl::cat(imageClassifierCat));
76} // unnamed namespace
77
78/// A pair representing a float and the index where the float was found.
79using FloatIndexPair = std::pair<float, size_t>;
80
81/// Given a Handle \p H of a 1D tensor with float elements, \returns the top K
82/// (topKCount) [float, index] pairs, i.e. the pairs with the highest floats.
83template <typename ElemTy>
84static std::vector<FloatIndexPair> getTopKPairs(Handle<ElemTy> H) {
85 DCHECK_LE(topKCount, H.size()) << "Function requires k < number of labels.";
86 DCHECK_EQ(H.dims().size(), 1) << "H must be a Handle of a 1d Tensor.";
87
88 // Use a priority queue of pairs of floats (probabilities) to size_t (indices)
89 // to determine the top K pairs, and then return the indices from it.
90 std::priority_queue<FloatIndexPair, std::vector<FloatIndexPair>,
91 std::greater<FloatIndexPair>>
92 topKQueue;
93
94 // Loop over all the probabilites, finding the highest k probability pairs.
95 for (dim_t i = 0, e = H.size(); i < e; i++) {
96 float currProbability = H.at({i});
97 if (topKQueue.size() < topKCount) {
98 // Always push the first k elements.
99 topKQueue.push(std::make_pair(currProbability, i));
100 } else if (topKQueue.top().first < currProbability) {
101 // If the lowest element has lower probability than the current, then pop
102 // the lowest and insert the current pair.
103 topKQueue.pop();
104 topKQueue.push(std::make_pair(currProbability, i));
105 }
106 }
107
108 // We now have the top K pairs in reverse order.
109 std::vector<FloatIndexPair> res(topKCount);
110 for (size_t i = 0; i < topKCount; i++) {
111 res[topKCount - i - 1] = topKQueue.top();
112 topKQueue.pop();
113 }
114
115 return res;
116}
117
118/// Print out the top K pairs to stdout, which were passed in via \p topKPairs.
119static void printTopKPairs(const std::vector<FloatIndexPair> &topKPairs) {
120 for (size_t i = 0; i < topKPairs.size(); i++) {
121 // Some models are trained with more classes. E.g. Some imagenet models
122 // exported from TensorFlow have 1 extra "neutral" class.
123 const size_t label = topKPairs[i].second - labelOffset;
124 // Tab out the label so it aligns nicely with Label-K1.
125 if (i != 0) {
126 llvm::outs() << "\t\t\t\t\t";
127 }
128 llvm::outs() << "\tLabel-K" << i + 1 << ": " << label << " (probability: "
129 << llvm::format("%0.4f", topKPairs[i].first) << ")\n";
130 }
131}
132
133/// Checks if \p topKPairs have the index that matches the provided index,
134/// \returns 0 on success and 1 if mismatches found.
135static int checkExpectedLabel(llvm::ArrayRef<FloatIndexPair> topKPairs,
136 llvm::StringRef fileName,
137 unsigned expectedCategoryIndex) {
138 // Loop through pairs and try to find a matching label.
139 for (const auto &p : topKPairs) {
140 if (p.second - labelOffset == expectedCategoryIndex) {
141 return 0;
142 }
143 }
144
145 llvm::outs() << " File: " << fileName
146 << " doesn't match index: " << expectedCategoryIndex
147 << " in the top " << topKPairs.size() << " pairs\n";
148
149 return 1;
150}
151
152/// Apply the softmax function to the given handle.
153template <typename ElemTy> static void applySoftmax(Handle<ElemTy> H) {
154 DCHECK_EQ(H.dims().size(), 1) << "H must be a Handle of a 1d Tensor.";
155 float denominator = 0.0f;
156
157 for (auto elem : H) {
158 denominator += std::exp(static_cast<float>(elem));
159 }
160
161 for (auto &elem : H) {
162 elem = std::exp(static_cast<float>(elem)) / denominator;
163 }
164}
165
166/// Given the output Softmax Tensor \p SMT and \p imageList, prints the
167/// results of inference and returns number of incorrect predictions,
168/// \returns the number of found mismatches.
169template <typename ElemTy>
170static int processAndPrintResultsImpl(Tensor *SMT,
171 llvm::ArrayRef<std::string> imageList) {
172 // Softmax should have at least two dimensions: batchSize (first dimension),
173 // numLabels (any other dimension), and optionally - 1 in all other
174 // dimensions. The value of numLabels should be greater than 1.
175 DCHECK_GE(SMT->dims().size(), 2) << "Softmax should have at least 2 dims.";
176 const dim_t batchSize = SMT->dims()[0];
177 DCHECK_EQ(batchSize, imageList.size())
178 << "Softmax batch size must equal the input number of images.";
179 size_t labelsDim = 0;
180 for (size_t i = 1; i < SMT->dims().size(); i++) {
181 if (SMT->dims()[i] > 1) {
182 DCHECK_EQ(labelsDim, 0) << "More than one dimension of size > 1?";
183 labelsDim = i;
184 }
185 }
186 DCHECK_NE(labelsDim, 0) << "Labels dimension not found!";
187 const dim_t numLabels = SMT->dims()[labelsDim];
188 // Get a view with canonical layout {batches, labels}.
189 Tensor canonical = SMT->getUnowned({batchSize, numLabels});
190 SMT = &canonical;
191
192 std::vector<dim_t> sliceOffset(SMT->dims().size(), 0);
193
194 int retVal = 0;
195 for (unsigned i = 0; i < imageList.size(); i++) {
196 const auto &fileName = imageList[i];
197 if (topKCount) {
198 llvm::outs() << " File: " << fileName;
199 }
200
201 // batchSize is the first dimension, so update it to get the next slice.
202 sliceOffset[0] = i;
203 Tensor slice = SMT->getUnowned({numLabels}, sliceOffset);
204 auto SH = slice.getHandle<ElemTy>();
205
206 if (computeSoftmax) {
207 applySoftmax(SH);
208 }
209
210 if (topKCount) {
211 auto topKPairs = getTopKPairs(SH);
212 printTopKPairs(topKPairs);
213 if (!expectedMatchingLabels.empty()) {
214 retVal +=
215 checkExpectedLabel(topKPairs, fileName, expectedMatchingLabels[i]);
216 }
217 }
218 }
219
220 return retVal;
221}
222
223class ImageClassifierProcessResult : public PostProcessOutputDataExtension {
224public:
225 int processOutputs(const llvm::StringMap<Placeholder *> &PHM,
226 PlaceholderBindings &bindings,
227 VecVecRef<std::string> imageList) override;
228};
229
230/// Given the output PlaceHolder StringMap \p PHM, of size 1, from SoftMax and
231/// \p functionName, switch between the correct element type to print the
232/// results of inference as contained in \p PHM, \returns the number of found
233/// mismatches.
234int ImageClassifierProcessResult::processOutputs(
235 const llvm::StringMap<Placeholder *> &PHM, PlaceholderBindings &bindings,
236 VecVecRef<std::string> imageList) {
237
238 if (profilingGraph()) {
239 LOG(INFO) << "Graph profiling is ON. Processing of output is disabled.";
240 return 0;
241 }
242
243 Placeholder *phOut = getOutputForPostProcessing(PHM);
244 if (!phOut) {
245 return 0;
246 }
247
248 auto *SMT = bindings.get(PHM.begin()->second);
249 switch (SMT->getElementType()) {
250 case ElemKind::FloatTy:
251 return processAndPrintResultsImpl<float>(SMT, imageList[0]);
252 case ElemKind::Float16Ty:
253 return processAndPrintResultsImpl<float16_t>(SMT, imageList[0]);
254 case ElemKind::BFloat16Ty:
255 return processAndPrintResultsImpl<bfloat16_t>(SMT, imageList[0]);
256 default:
257 llvm_unreachable("Type not supported");
258 }
259 return 0;
260}
261
262int main(int argc, char **argv) {
263 google::InitGoogleLogging(argv[0]);
264
265 if (!expectedMatchingLabels.empty()) {
266 // The number of category indices must match the number of files.
267 if (expectedMatchingLabels.size() != inputImageFilenamesOpt[0].size()) {
268 llvm::errs() << "Number of matching indices: "
269 << expectedMatchingLabels.size()
270 << " doesn't match the number of files: "
271 << inputImageFilenamesOpt[0].size() << "\n";
272 return 1;
273 }
274 }
275
276 glow::Executor core("ImageClassifier", argc, argv);
277 auto printResultCreator =
278 []() -> std::unique_ptr<PostProcessOutputDataExtension> {
279 return std::make_unique<ImageClassifierProcessResult>();
280 };
281 core.registerPostProcessOutputExtension(printResultCreator);
282
283 int numErrors = core.executeNetwork();
284 return numErrors;
285}
286