pytorch

1
#include <gtest/gtest.h>
2
#include "caffe2/core/net.h"
3
#include "caffe2/core/operator.h"
4
#include "caffe2/transforms/conv_to_nnpack_transform.h"
5

6
namespace caffe2 {
7

8
namespace {
9

10
using transform::Graph;
11

12
TEST(ConvToNNPackTest, TestSimple) {
13
  NetDef netdef;
14
  // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
15
  OperatorDef* op;
16
  // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
17
  op = AddOp(&netdef, "Conv", {"in"}, {"out"});
18
  // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
19
  op = AddOp(&netdef, "Relu", {"out"}, {"out"});
20
  op = AddOp(&netdef, "Conv", {"out"}, {"out"}); // if not CPU, won't transform
21
  op->mutable_device_option()->set_device_type(PROTO_CUDA);
22
  // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
23
  op = AddOp(&netdef, "Relu", {"out"}, {"out"});
24
  op = AddOp(&netdef, "Conv", {"out"}, {"out"});
25
  op->set_engine("NNPACK"); // does not need to be transformed
26
  // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
27
  op = AddOp(&netdef, "Relu", {"out"}, {"out"});
28
  // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
29
  op = AddOp(&netdef, "Conv", {"out"}, {"out"});
30
  // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
31
  op = AddOp(&netdef, "Relu", {"out"}, {"out"});
32

33
  auto t = TransformRegistry()->Create("ConvToNNPack");
34
  NetDef transformed_netdef = t->ApplyTo(netdef);
35

36
  int nnpack_count = 0;
37
  for (auto& op : transformed_netdef.op()) {
38
    if (op.type() == "Conv" && op.device_option().device_type() == PROTO_CPU) {
39
      EXPECT_EQ(op.engine(), "NNPACK");
40
      nnpack_count++;
41
    }
42
  }
43
  EXPECT_EQ(nnpack_count, 3);
44
  EXPECT_EQ(t->PatternMatch(Graph(netdef)).size(), 2); // should get 2 matches
45
}
46

47
} // namespace
48

49
} // namespace caffe2
50