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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
|
#include <iostream>
#include <memory>
#include "caffe2/core/module.h"
#include "caffe2/core/operator.h"
#include <gtest/gtest.h>
#include "caffe2/core/logging.h"
// An explicitly defined module, testing correctness when we statically link a
// module
CAFFE2_MODULE(caffe2_module_test_static, "Static module for testing.");
namespace caffe2 {
class Caffe2ModuleTestStaticDummyOp : public OperatorBase {
public:
using OperatorBase::OperatorBase;
bool Run(int /* unused */ /*stream_id*/) override {
return true;
}
virtual string type() {
return "base";
}
};
REGISTER_CPU_OPERATOR(
Caffe2ModuleTestStaticDummy, Caffe2ModuleTestStaticDummyOp);
OPERATOR_SCHEMA(Caffe2ModuleTestStaticDummy);
TEST(ModuleTest, StaticModule) {
const string name = "caffe2_module_test_static";
const auto& modules = CurrentModules();
EXPECT_EQ(modules.count(name), 1);
EXPECT_TRUE(HasModule(name));
// LoadModule should not raise an error, since the module is already present.
LoadModule(name);
// Even a non-existing path should not cause error.
LoadModule(name, "/does/not/exist.so");
EXPECT_EQ(modules.count(name), 1);
EXPECT_TRUE(HasModule(name));
// The module will then introduce the Caffe2ModuleTestStaticDummyOp.
OperatorDef op_def;
Workspace ws;
op_def.set_type("Caffe2ModuleTestStaticDummy");
unique_ptr<OperatorBase> op = CreateOperator(op_def, &ws);
EXPECT_NE(nullptr, op.get());
}
#ifdef CAFFE2_BUILD_SHARED_LIBS
TEST(ModuleTest, DynamicModule) {
const string name = "caffe2_module_test_dynamic";
const auto& modules = CurrentModules();
EXPECT_EQ(modules.count(name), 0);
EXPECT_FALSE(HasModule(name));
// Before loading, we should not be able to create the op.
OperatorDef op_def;
Workspace ws;
op_def.set_type("Caffe2ModuleTestDynamicDummy");
EXPECT_THROW(
CreateOperator(op_def, &ws),
EnforceNotMet);
// LoadModule should load the proper module.
LoadModule(name);
EXPECT_EQ(modules.count(name), 1);
EXPECT_TRUE(HasModule(name));
// The module will then introduce the Caffe2ModuleTestDynamicDummyOp.
unique_ptr<OperatorBase> op_after_load = CreateOperator(op_def, &ws);
EXPECT_NE(nullptr, op_after_load.get());
}
#endif
} // namespace caffe2
|
