Caffe2 - C++ API
A deep learning, cross platform ML framework
conv_to_nnpack_transform.h
1 
17 #pragma once
18 
19 #include "caffe2/core/common.h"
20 #include "caffe2/proto/caffe2.pb.h"
21 #include "caffe2/transforms/single_op_transform.h"
22 #include "caffe2/utils/proto_utils.h"
23 
24 namespace caffe2 {
25 
27  protected:
28  // Specify what the op needs to be to match the pattern.
29  bool MatchOperator(const OperatorDef& op) override {
30  return (
31  op.type() == "Conv" && op.device_option().device_type() == CPU &&
32  op.engine() != "NNPACK");
33  }
34 
35  // Specify how the operator should be replaced.
36  void ReplaceOperator(OperatorDef* op) override {
37  op->set_engine("NNPACK");
38  }
39 };
40 
41 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.
Single Op Transform Base class.