pytorch  1.8.2
About: PyTorch provides Tensor computation (like NumPy) with strong GPU acceleration and Deep Neural Networks (in Python) built on a tape-based autograd system. LTS (Long Term Support) release.
  Fossies Dox: pytorch-1.8.2.tar.gz  ("unofficial" and yet experimental doxygen-generated source code documentation)  
operator_schema_test.cc
Go to the documentation of this file.
1#include "caffe2/core/logging.h"
2#include "caffe2/core/operator.h"
3#include "caffe2/core/operator_schema.h"
4#include "caffe2/utils/proto_utils.h"
5
6#include <gtest/gtest.h>
7
8namespace caffe2 {
9
10OPERATOR_SCHEMA(OpSchemaTestOp)
11 .NumInputs(1).NumOutputs(1)
12 .SetDoc(R"DOC(Test Documentation)DOC")
13 .Input(0, "in0", "dummy input.")
14 .Output(0, "out0", "dummy output.");
15
16TEST(OperatorSchemaTest, BasicSchema) {
17 const OpSchema* schema = OpSchemaRegistry::Schema("OpSchemaTestOp");
18#ifdef CAFFE2_NO_OPERATOR_SCHEMA
19 EXPECT_TRUE(schema == nullptr);
20 return;
21#endif
22 EXPECT_TRUE(schema != nullptr);
23 EXPECT_TRUE(schema->doc() != nullptr);
24 OperatorDef def1 = CreateOperatorDef(
25 "OpSchemaTestOp", "",
26 vector<string>{"in"}, vector<string>{"out"});
27 EXPECT_TRUE(schema->Verify(def1));
28 OperatorDef def2 = CreateOperatorDef(
29 "OpSchemaTestOp", "",
30 vector<string>{"in1", "in2"}, vector<string>{"out"});
31 EXPECT_FALSE(schema->Verify(def2));
32 OperatorDef def3 = CreateOperatorDef(
33 "OpSchemaTestOp", "",
34 vector<string>{"in"}, vector<string>{"out1", "out2"});
35 EXPECT_FALSE(schema->Verify(def3));
36}
37
38OPERATOR_SCHEMA(OpSchemaSpecifiedInputOutputOp)
39 .NumInputs({2, 4}).NumOutputs({1, 3});
40
41TEST(OperatorSchemaTest, SpecifiedInputOutput) {
42 const OpSchema* schema
43 = OpSchemaRegistry::Schema("OpSchemaSpecifiedInputOutputOp");
44#ifdef CAFFE2_NO_OPERATOR_SCHEMA
45 EXPECT_TRUE(schema == nullptr);
46 return;
47#endif
48 EXPECT_TRUE(schema != nullptr);
49 OperatorDef def1 = CreateOperatorDef(
50 "OpSchemaSpecifiedInputOutputOp", "",
51 vector<string>{"in"}, vector<string>{"out"});
52 EXPECT_FALSE(schema->Verify(def1));
53 OperatorDef def2 = CreateOperatorDef(
54 "OpSchemaSpecifiedInputOutputOp", "",
55 vector<string>{"in1", "in2"}, vector<string>{"out"});
56 EXPECT_TRUE(schema->Verify(def2));
57 OperatorDef def3 = CreateOperatorDef(
58 "OpSchemaSpecifiedInputOutputOp", "",
59 vector<string>{"in1", "in2"}, vector<string>{"out1", "out2"});
60 EXPECT_FALSE(schema->Verify(def3));
61}
62
63OPERATOR_SCHEMA(OpSchemaInputOutputRelationOp)
64 .NumInputsOutputs([](int in, int out) {
65 return out == in || out == in * 2;
66 });
67
68TEST(OperatorSchemaTest, InputOutputRelation) {
69 const OpSchema* schema
70 = OpSchemaRegistry::Schema("OpSchemaInputOutputRelationOp");
71#ifdef CAFFE2_NO_OPERATOR_SCHEMA
72 EXPECT_TRUE(schema == nullptr);
73 return;
74#endif
75 EXPECT_TRUE(schema != nullptr);
76 OperatorDef def1 = CreateOperatorDef(
77 "OpSchemaInputOutputRelationOp", "",
78 vector<string>{"in"}, vector<string>{"out"});
79 EXPECT_TRUE(schema->Verify(def1));
80 OperatorDef def2 = CreateOperatorDef(
81 "OpSchemaInputOutputRelationOp", "",
82 vector<string>{"in"}, vector<string>{"out1", "out2"});
83 EXPECT_TRUE(schema->Verify(def2));
84 OperatorDef def3 = CreateOperatorDef(
85 "OpSchemaInputOutputRelationOp", "",
86 vector<string>{"in1", "in2", "in3"}, vector<string>{"out1", "out2"});
87 EXPECT_FALSE(schema->Verify(def3));
88}
89
90OPERATOR_SCHEMA(OpSchemaSameInputOutputOp)
91 .SameNumberOfOutput();
92
93TEST(OperatorSchemaTest, SameInputOutput) {
94 const OpSchema* schema =
95 OpSchemaRegistry::Schema("OpSchemaSameInputOutputOp");
96#ifdef CAFFE2_NO_OPERATOR_SCHEMA
97 EXPECT_TRUE(schema == nullptr);
98 return;
99#endif
100 OperatorDef def1 = CreateOperatorDef(
101 "OpSchemaSameInputOutputOp", "",
102 vector<string>{"in"}, vector<string>{"out"});
103 EXPECT_TRUE(schema->Verify(def1));
104 OperatorDef def2 = CreateOperatorDef(
105 "OpSchemaSameInputOutputOp", "",
106 vector<string>{"in1", "in2"}, vector<string>{"out1", "out2"});
107 EXPECT_TRUE(schema->Verify(def2));
108 OperatorDef def3 = CreateOperatorDef(
109 "OpSchemaSameInputOutputOp", "",
110 vector<string>{"in1", "in2"}, vector<string>{"out1", "out2", "out3"});
111 EXPECT_FALSE(schema->Verify(def3));
112}
113
114OPERATOR_SCHEMA(OpSchemaCalculateOutputOp)
115 .NumInputs(1, 5).NumOutputs(2, 6)
116 .OutputCalculator([](int n) { return n + 1; });
117
118TEST(OperatorSchemaTest, CalculateOutput) {
119 const OpSchema* schema =
120 OpSchemaRegistry::Schema("OpSchemaCalculateOutputOp");
121#ifdef CAFFE2_NO_OPERATOR_SCHEMA
122 EXPECT_TRUE(schema == nullptr);
123 return;
124#endif
125 OperatorDef def1 = CreateOperatorDef(
126 "OpSchemaCalculateOutputOp", "",
127 vector<string>{"in"}, vector<string>{"out"});
128 EXPECT_FALSE(schema->Verify(def1));
129 OperatorDef def2 = CreateOperatorDef(
130 "OpSchemaCalculateOutputOp", "",
131 vector<string>{"in1", "in2"}, vector<string>{"out1", "out2"});
132 EXPECT_FALSE(schema->Verify(def2));
133 OperatorDef def3 = CreateOperatorDef(
134 "OpSchemaCalculateOutputOp", "",
135 vector<string>{"in1", "in2"}, vector<string>{"out1", "out2", "out3"});
136 EXPECT_TRUE(schema->Verify(def3));
137}
138
139OPERATOR_SCHEMA(OpSchemaInplace)
140 .NumInputs(2).NumOutputs(2)
141 .AllowInplace({{0, 0}})
142 .EnforceInplace({{1, 1}});
143
144TEST(OperatorSchemaTest, Inplace) {
145 const OpSchema* schema =
146 OpSchemaRegistry::Schema("OpSchemaInplace");
147#ifdef CAFFE2_NO_OPERATOR_SCHEMA
148 EXPECT_TRUE(schema == nullptr);
149 return;
150#endif
151 OperatorDef def1 = CreateOperatorDef(
152 "OpSchemaInplace", "",
153 vector<string>{"in1", "in2"}, vector<string>{"out1", "in2"});
154 EXPECT_TRUE(schema->Verify(def1));
155 OperatorDef def2 = CreateOperatorDef(
156 "OpSchemaInplace", "",
157 vector<string>{"in1", "in2"}, vector<string>{"in1", "in2"});
158 EXPECT_TRUE(schema->Verify(def2));
159 OperatorDef def3 = CreateOperatorDef(
160 "OpSchemaInplace", "",
161 vector<string>{"in1", "in2"}, vector<string>{"in1", "out2"});
162 EXPECT_FALSE(schema->Verify(def3));
163 OperatorDef def4 = CreateOperatorDef(
164 "OpSchemaInplace", "",
165 vector<string>{"in1", "in2"}, vector<string>{"out1", "out2"});
166 EXPECT_FALSE(schema->Verify(def4));
167}
168
169OPERATOR_SCHEMA(OpSchemaSameInputOutputTensorInference).IdenticalTypeAndShape();
170
171TEST(OperatorSchemaTest, TensorInferenceIdentical) {
172 const OpSchema* schema =
173 OpSchemaRegistry::Schema("OpSchemaSameInputOutputTensorInference");
174#ifdef CAFFE2_NO_OPERATOR_SCHEMA
175 EXPECT_TRUE(schema == nullptr);
176 return;
177#endif
178 OperatorDef def = CreateOperatorDef(
179 "OpSchemaSameInputOutputTensorInference",
180 "",
181 vector<string>{"in"},
182 vector<string>{"out"});
183 vector<TensorShape> shapes(1);
184 shapes[0].set_data_type(TensorProto::FLOAT);
185 shapes[0].add_dims(1);
186 shapes[0].add_dims(2);
187 shapes[0].add_dims(3);
188 vector<TensorShape> out = schema->InferTensor(def, shapes);
189 EXPECT_EQ(out.size(), 1);
190 EXPECT_EQ(out[0].SerializeAsString(), shapes[0].SerializeAsString());
191}
192
193OPERATOR_SCHEMA(OpSchemaArbitraryTensorInference)
194 .TensorInferenceFunction(
195 [](const OperatorDef&, const vector<TensorShape>&) {
196 vector<TensorShape> shapes(1);
197 shapes[0].set_data_type(TensorProto::FLOAT);
198 shapes[0].add_dims(1701);
199 return shapes;
200 });
201
202TEST(OperatorSchemaTest, TensorInferenceArbitrary) {
203 const OpSchema* schema =
204 OpSchemaRegistry::Schema("OpSchemaArbitraryTensorInference");
205#ifdef CAFFE2_NO_OPERATOR_SCHEMA
206 EXPECT_TRUE(schema == nullptr);
207 return;
208#endif
209 OperatorDef def = CreateOperatorDef(
210 "OpSchemaArbitraryTensorInference",
211 "",
212 vector<string>{"in"},
213 vector<string>{"out"});
214 vector<TensorShape> out = schema->InferTensor(def, vector<TensorShape>());
215 EXPECT_EQ(out.size(), 1);
216 EXPECT_EQ(out[0].data_type(), TensorProto::FLOAT);
217 EXPECT_EQ(out[0].dims_size(), 1);
218 EXPECT_EQ(out[0].dims(0), 1701);
219}
220
221TEST(OperatorSchemaTest, TestCastSchema) {
222 // This tests a use case of the schema: the Cast op takes in the def and
223 // deduces the
224 // schema from the "to" argument.
225 const OpSchema* schema = OpSchemaRegistry::Schema("Cast");
226#ifdef CAFFE2_NO_OPERATOR_SCHEMA
227 EXPECT_TRUE(schema == nullptr);
228 return;
229#endif
230 if (!schema) {
231 // Compiled without the Cast op.
232 return;
233 }
234 OperatorDef def = CreateOperatorDef(
235 "Cast",
236 "",
237 vector<string>{"in"},
238 vector<string>{"out"},
239 vector<Argument>{MakeArgument<int64_t>("to", TensorProto::UINT8)});
240 vector<TensorShape> out = schema->InferTensor(def, vector<TensorShape>(1));
241 EXPECT_EQ(out.size(), 1);
242 // Data type should be inferred.
243 EXPECT_EQ(out[0].data_type(), TensorProto::UINT8);
244 // Dim should not be set (same as input);
245 EXPECT_EQ(out[0].dims_size(), 0);
246}
247
248OPERATOR_SCHEMA(OpSchemaCostInference)
249 .NumInputs(2)
250 .NumOutputs(2)
251 .CostInferenceFunction([](const OperatorDef& /*def*/,
252 const vector<TensorShape>& inputs) {
253 struct OpSchema::Cost c;
254 c.flops = 2 * inputs[0].dims(0) * inputs[0].dims(1) * inputs[1].dims(1);
255 return c;
256 });
257
258TEST(OperatorSchemaTest, TestCostInference) {
259 const OpSchema* schema = OpSchemaRegistry::Schema("OpSchemaCostInference");
260#ifdef CAFFE2_NO_OPERATOR_SCHEMA
261 EXPECT_TRUE(schema == nullptr);
262 return;
263#endif
264 if (!schema) {
265 return;
266 }
267 OperatorDef def = CreateOperatorDef(
268 "OpSchemaCostInference", "", vector<string>{"in"}, vector<string>{"out"});
269 vector<TensorShape> shapes(2);
270 shapes[0].set_data_type(TensorProto::FLOAT);
271 shapes[0].add_dims(10);
272 shapes[0].add_dims(10);
273 shapes[1].set_data_type(TensorProto::FLOAT);
274 shapes[1].add_dims(10);
275 shapes[1].add_dims(10);
276 EXPECT_EQ(2000, schema->InferCost(def, shapes).flops);
277}
278
279} // namespace caffe2
caffe2::OpSchemaRegistry::Schema
static const OpSchema * Schema(const string &key)
Definition: operator_schema.h:479
caffe2::OpSchema
A class to record the schema of an op.
Definition: operator_schema.h:40
caffe2
Copyright (c) 2016-present, Facebook, Inc.
Definition: blob.h:13
caffe2::inputs
const vector< TensorShape > & inputs
Definition: operator_schema_test.cc:252
caffe2::OPERATOR_SCHEMA
OPERATOR_SCHEMA(ATen)
caffe2::EnforceInplace
EnforceInplace({{1, 1}})
caffe2::vector< TensorShape >
return vector< TensorShape >
Definition: slice_op.cc:110
caffe2::in
const vector< TensorShape > & in
Definition: batch_gather_ops.cc:12
caffe2::NumOutputs
INT_MAX NumOutputs(1, INT_MAX)
caffe2::dims
*and produces a single output tensor *expanded *The op also takes an argument *dims *with a list of dimensions for where to add the single dimensional entries If the same blob is provided as input and the operation is copy free This is the exact inverse operation of *Squeeze *Github dims
Definition: expand_squeeze_dims_op.cc:72
caffe2::CreateOperatorDef
OperatorDef CreateOperatorDef(const string &type, const string &name, const IterableInputs &inputs, const IterableOutputs &outputs, const IterableArgs &args, const DeviceOption &device_option=DeviceOption(), const string &engine="")
Definition: proto_utils.h:141
caffe2::TEST
TEST(CommonTest, TestStoi)
Definition: common_test.cc:16
schema
Module caffe2.python.schema.
operator_schema.h
proto_utils.h