forked from chromium/chromium
-
Notifications
You must be signed in to change notification settings - Fork 0
/
nn_classifier.cc
95 lines (77 loc) · 2.73 KB
/
nn_classifier.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
// Copyright (c) 2018 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "components/assist_ranker/nn_classifier.h"
#include "base/check_op.h"
#include "components/assist_ranker/proto/nn_classifier.pb.h"
namespace assist_ranker {
namespace nn_classifier {
namespace {
using google::protobuf::RepeatedPtrField;
using std::vector;
vector<float> FeedForward(const NNLayer& layer, const vector<float>& input) {
const RepeatedPtrField<FloatVector>& weights = layer.weights();
const FloatVector& biases = layer.biases();
// Number of nodes in the layer.
const int num_nodes = biases.values().size();
// Number of values in the input.
const int num_input = input.size();
DCHECK_EQ(weights.size(), num_input);
// Initialize with the bias.
vector<float> output(biases.values().begin(), biases.values().end());
// For each value in the input.
for (int j = 0; j < num_input; ++j) {
const FloatVector& v = weights[j];
DCHECK_EQ(v.values().size(), num_nodes);
// For each node in the layer.
for (int i = 0; i < num_nodes; ++i) {
output[i] += v.values(i) * input[j];
}
}
return output;
}
// Apply ReLU activation function to a vector, which sets all values to
// max(0, value).
void Relu(vector<float>* const v) {
// We are modifying the vector so the iterator must be a reference.
for (float& i : *v)
if (i < 0.0f)
i = 0.0f;
}
bool ValidateLayer(const NNLayer& layer) {
// Number of nodes in the layer (must be non-zero).
const int num_nodes = layer.biases().values().size();
if (num_nodes == 0)
return false;
// Number of values in the input (must be non-zero).
const int num_input = layer.weights().size();
if (num_input == 0)
return false;
for (int j = 0; j < num_input; ++j) {
// The size of each weight vector must be the number of nodes in the
// layer.
if (layer.weights(j).values().size() != num_nodes)
return false;
}
return true;
}
} // namespace
bool Validate(const NNClassifierModel& model) {
// Check the size of the output from the hidden layer is equal to the size
// of the input in the logits layer.
if (model.hidden_layer().biases().values().size() !=
model.logits_layer().weights().size()) {
return false;
}
return ValidateLayer(model.hidden_layer()) &&
ValidateLayer(model.logits_layer());
}
vector<float> Inference(const NNClassifierModel& model,
const vector<float>& input) {
vector<float> v = FeedForward(model.hidden_layer(), input);
Relu(&v);
// Feed forward the logits layer.
return FeedForward(model.logits_layer(), v);
}
} // namespace nn_classifier
} // namespace assist_ranker