forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
glow_net_transform.h
40 lines (34 loc) · 1.25 KB
/
glow_net_transform.h
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
#pragma once
#include <string>
#include <unordered_set>
#include <vector>
#include <caffe2/core/net.h>
#include <caffe2/core/workspace.h>
#include <caffe2/opt/shape_info.h>
#include <caffe2/proto/caffe2_pb.h>
C10_DECLARE_string(onnxifi_blacklist);
C10_DECLARE_string(onnxifi_blacklist_ops);
namespace caffe2 {
namespace glow {
/// Onnxifi transformation on the net and workspace. We also
/// needed the input data/shape to populate the shape. In addition, we take a \p
/// blacklist to control and mask what ops we want to consider in onnxifi
/// process. We can also set whether to use ONNX proto or C2 proto through
/// ONNXIFI interface.
void onnxifi(
NetDef* net,
Workspace* ws,
const std::vector<std::string>& input_names,
const std::vector<std::string>& output_names,
const std::vector<std::string>& weight_names,
const std::unordered_set<int>& blacklist,
const ShapeInfoMap& shape_hints,
bool use_onnx,
size_t max_batch_size = 0,
size_t max_seq_size = 0,
bool load_model_by_blob = false,
bool predictor_net_ssa_rewritten = false);
std::unordered_set<int> ParseNetPositionList(const std::string& str);
std::unordered_set<std::string> ParseBlackListOps(const std::string& str);
} // namespace glow
} // namespace caffe2