1 #include <gtest/gtest.h>
2 #include <test/cpp/jit/test_utils.h>
3 #include <torch/torch.h>
4 #include <torch/csrc/jit/api/module.h>
5 #include <torch/csrc/jit/frontend/resolver.h>
6 #include <torch/csrc/jit/mobile/import.h>
7 #include <torch/csrc/jit/mobile/module.h>
8 // Cover codegen'd unboxing logic for these types:
9 //'Device',
10 //'Device?',
11 //'Dimname',
12 //'Dimname[1]',
13 //'Dimname[]',
14 //'Dimname[]?',
15 //'Generator?',
16 //'Layout?',
17 //'MemoryFormat',
18 //'MemoryFormat?',
19 //'Scalar',
20 //'Scalar?',
21 //'ScalarType',
22 //'ScalarType?',
23 //'Scalar[]',
24 //'Storage',
25 //'Stream',
26 //'Tensor',
27 //'Tensor(a!)',
28 //'Tensor(a!)[]',
29 //'Tensor(a)',
30 //'Tensor(b!)',
31 //'Tensor(c!)',
32 //'Tensor(d!)',
33 //'Tensor?',
34 //'Tensor?[]',
35 //'Tensor[]',
36 //'bool',
37 //'bool?',
38 //'bool[2]',
39 //'bool[3]',
40 //'bool[4]',
41 //'float',
42 //'float?',
43 //'float[]?',
44 //'int',
45 //'int?',
46 //'int[1]',
47 //'int[1]?',
48 //'int[2]',
49 //'int[2]?',
50 //'int[3]',
51 //'int[4]',
52 //'int[5]',
53 //'int[6]',
54 //'int[]',
55 //'int[]?',
56 //'str',
57 //'str?'
58 namespace torch {
59 namespace jit {
60 namespace mobile {
61 // covers int[], ScalarType?, Layout?, Device?, bool?
TEST(LiteInterpreterTest,Ones)62 TEST(LiteInterpreterTest, Ones) {
63 // Load check in model: ModelWithDTypeDeviceLayoutPinMemory.ptl
64 auto testModelFile = "ModelWithDTypeDeviceLayoutPinMemory.ptl";
65
66 // class ModelWithDTypeDeviceLayoutPinMemory(torch.nn.Module):
67 // def forward(self, x: int):
68 // a = torch.ones([3, x], dtype=torch.int64, layout=torch.strided, device="cpu")
69 // return a
70 Module bc = _load_for_mobile(testModelFile);
71 std::vector<c10::IValue> input{c10::IValue(4)};
72 const auto result = bc.forward(input);
73 ASSERT_EQ(result.toTensor().size(0), 3);
74 ASSERT_EQ(result.toTensor().size(1), 4);
75 }
76
TEST(LiteInterpreterTest,Index)77 TEST(LiteInterpreterTest, Index) {
78 // Load check in model: ModelWithTensorOptional.ptl
79 auto testModelFile = "ModelWithTensorOptional.ptl";
80
81 // class ModelWithTensorOptional(torch.nn.Module):
82 // def forward(self, index):
83 // a = torch.zeros(2, 2)
84 // a[0][1] = 1
85 // a[1][0] = 2
86 // a[1][1] = 3
87 // return a[index]
88 Module bc = _load_for_mobile(testModelFile);
89 int64_t ind_1 = 0;
90
91 const auto result_1 = bc.forward({at::tensor(ind_1)});
92
93 at::Tensor expected = at::empty({1, 2}, c10::TensorOptions(c10::ScalarType::Float));
94 expected[0][0] = 0;
95 expected[0][1] = 1;
96
97 AT_ASSERT(result_1.toTensor().equal(expected));
98 }
99
TEST(LiteInterpreterTest,Gradient)100 TEST(LiteInterpreterTest, Gradient) {
101 // Load check in model: ModelWithScalarList.ptl
102 auto testModelFile = "ModelWithScalarList.ptl";
103
104 // class ModelWithScalarList(torch.nn.Module):
105 // def forward(self, a: int):
106 // values = torch.tensor([4., 1., 1., 16.], )
107 // if a == 0:
108 // return torch.gradient(values, spacing=torch.scalar_tensor(2., dtype=torch.float64))
109 // elif a == 1:
110 // return torch.gradient(values, spacing=[torch.tensor(1.).item()])
111 Module bc = _load_for_mobile(testModelFile);
112
113 const auto result_1 = bc.forward({0});
114 at::Tensor expected_1 = at::tensor({-1.5, -0.75, 3.75, 7.5}, c10::TensorOptions(c10::ScalarType::Float));
115 AT_ASSERT(result_1.toList().get(0).toTensor().equal(expected_1));
116
117 const auto result_2 = bc.forward({1});
118 at::Tensor expected_2 = at::tensor({-3.0, -1.5, 7.5, 15.0}, c10::TensorOptions(c10::ScalarType::Float));
119 AT_ASSERT(result_2.toList().get(0).toTensor().equal(expected_2));
120 }
121
TEST(LiteInterpreterTest,Upsample)122 TEST(LiteInterpreterTest, Upsample) {
123 // Load check in model: ModelWithFloatList.ptl
124 auto testModelFile = "ModelWithFloatList.ptl";
125
126 // model = torch.nn.Upsample(scale_factor=(2.0,), mode="linear")
127 Module bc = _load_for_mobile(testModelFile);
128
129 const auto result_1 = bc.forward({at::ones({1, 2, 3})});
130 at::Tensor expected_1 = at::ones({1, 2, 6}, c10::TensorOptions(c10::ScalarType::Float));
131 AT_ASSERT(result_1.toTensor().equal(expected_1));
132 }
133
TEST(LiteInterpreterTest,IndexTensor)134 TEST(LiteInterpreterTest, IndexTensor) {
135 // Load check in model: ModelWithListOfOptionalTensors.ptl
136 auto testModelFile = "ModelWithListOfOptionalTensors.ptl";
137
138 // class ModelWithListOfOptionalTensors(torch.nn.Module):
139 // def forward(self, index):
140 // values = torch.tensor([4., 1., 1., 16.], )
141 // return values[[index, torch.tensor(0)]]
142 Module bc = _load_for_mobile(testModelFile);
143 const auto result_1 = bc.forward({at::tensor({1}, c10::TensorOptions(c10::ScalarType::Long))});
144
145 at::Tensor expected_1 = at::tensor({1.}, c10::TensorOptions(c10::ScalarType::Float));
146 AT_ASSERT(result_1.toTensor().equal(expected_1));
147 }
148
TEST(LiteInterpreterTest,Conv2d)149 TEST(LiteInterpreterTest, Conv2d) {
150 // Load check in model: ModelWithArrayOfInt.ptl
151 auto testModelFile = "ModelWithArrayOfInt.ptl";
152
153 // model = torch.nn.Conv2d(1, 2, (2, 2), stride=(1, 1), padding=(1, 1))
154 Module bc = _load_for_mobile(testModelFile);
155 const auto result_1 = bc.forward({at::ones({1, 1, 1, 1})});
156
157 ASSERT_EQ(result_1.toTensor().sizes(), c10::IntArrayRef ({1,2,2,2}));
158 }
159
TEST(LiteInterpreterTest,AddTensor)160 TEST(LiteInterpreterTest, AddTensor) {
161 // Load check in model: ModelWithTensors.ptl
162 auto testModelFile = "ModelWithTensors.ptl";
163
164 // class ModelWithTensors(torch.nn.Module):
165 // def forward(self, a):
166 // values = torch.ones(size=[2, 3], names=['N', 'C'])
167 // values[0][0] = a[0]
168 // return values
169 Module bc = _load_for_mobile(testModelFile);
170 const auto result_1 = bc.forward({at::tensor({1, 2, 3}, c10::TensorOptions(c10::ScalarType::Long))});
171
172 at::Tensor expected_1 = at::tensor({2, 3, 4}, c10::TensorOptions(c10::ScalarType::Long));
173 AT_ASSERT(result_1.toTensor().equal(expected_1));
174 }
175
TEST(LiteInterpreterTest,DivideTensor)176 TEST(LiteInterpreterTest, DivideTensor) {
177 // Load check in model: ModelWithStringOptional.ptl
178 auto testModelFile = "ModelWithStringOptional.ptl";
179
180 // class ModelWithStringOptional(torch.nn.Module):
181 // def forward(self, b):
182 // a = torch.tensor(3, dtype=torch.int64)
183 // out = torch.empty(size=[1], dtype=torch.float)
184 // torch.div(b, a, out=out)
185 // return [torch.div(b, a, rounding_mode='trunc'), out]
186 Module bc = _load_for_mobile(testModelFile);
187 const auto result_1 = bc.forward({at::tensor({-12}, c10::TensorOptions(c10::ScalarType::Long))});
188
189 at::Tensor expected_1 = at::tensor({-4}, c10::TensorOptions(c10::ScalarType::Long));
190 at::Tensor expected_2 = at::tensor({-4.}, c10::TensorOptions(c10::ScalarType::Float));
191 AT_ASSERT(result_1.toList().get(0).toTensor().equal(expected_1));
192 AT_ASSERT(result_1.toList().get(1).toTensor().equal(expected_2));
193 }
194
TEST(LiteInterpreterTest,MultipleOps)195 TEST(LiteInterpreterTest, MultipleOps) {
196 // Load check in model: ModelWithMultipleOps.ptl
197 auto testModelFile = "ModelWithMultipleOps.ptl";
198
199 // class ModelWithMultipleOps(torch.nn.Module):
200 // def __init__(self) -> None:
201 // super().__init__()
202 // self.ops = torch.nn.Sequential(
203 // torch.nn.ReLU(),
204 // torch.nn.Flatten(),
205 // )
206 //
207 // def forward(self, x):
208 // x[1] = -2
209 // return self.ops(x)
210
211 Module bc = _load_for_mobile(testModelFile);
212 auto b = at::ones({2, 2, 2, 2});
213 const auto result = bc.forward({b});
214
215 at::Tensor expected = torch::tensor({{1, 1, 1, 1, 1, 1, 1, 1}, {0, 0, 0, 0, 0, 0, 0, 0}}, c10::TensorOptions(c10::ScalarType::Float));
216 AT_ASSERT(result.toTensor().equal(expected));
217 }
218 } // namespace mobile
219 } // namespace jit
220 } // namespace torch
221