pytorch

1
#include "caffe2/transforms/single_op_transform.h"
2

3
#include "caffe2/core/common.h"
4
#include "caffe2/core/logging.h"
5
#include "caffe2/core/net.h"
6
#include "caffe2/proto/caffe2_pb.h"
7

8
namespace caffe2 {
9

10
using transform::Graph;
11

12
bool SingleOpTransform::PatternRule(
13
    const Graph& g,
14
    const std::vector<int>& subgraph,
15
    int idx) {
16
  if (subgraph.size() == 0) {
17
    return MatchOperator(g.node(idx).op);
18
  }
19
  return false;
20
}
21

22
bool SingleOpTransform::ValidatorRule(
23
    const Graph& /*g*/,
24
    const std::vector<int>& subgraph) {
25
  if (subgraph.size() == 1) {
26
    return true;
27
  }
28
  return false;
29
}
30

31
bool SingleOpTransform::ReplaceRule(
32
    const std::vector<int>& subgraph,
33
    Graph* g_ptr) {
34
  CHECK(g_ptr);
35
  auto& g = *g_ptr;
36
  ReplaceOperator(&(g.node(subgraph[0]).op));
37
  return true;
38
}
39

40
} // namespace caffe2
41