Caffe2 - C++ API
A deep learning, cross platform ML framework
serialize.cpp
1 #include <gtest/gtest.h>
2 
3 #include <c10/util/tempfile.h>
4 
5 #include <torch/nn/modules/functional.h>
6 #include <torch/nn/modules/linear.h>
7 #include <torch/nn/modules/sequential.h>
8 #include <torch/optim/optimizer.h>
9 #include <torch/optim/sgd.h>
10 #include <torch/serialize.h>
11 #include <torch/types.h>
12 #include <torch/utils.h>
13 
14 #include <test/cpp/api/support.h>
15 
16 #include <cstdio>
17 #include <memory>
18 #include <sstream>
19 #include <string>
20 #include <vector>
21 
22 using namespace torch::nn;
23 using namespace torch::serialize;
24 
25 namespace {
26 Sequential xor_model() {
27  return Sequential(
28  Linear(2, 8),
29  Functional(at::sigmoid),
30  Linear(8, 1),
31  Functional(at::sigmoid));
32 }
33 
34 torch::Tensor save_and_load(torch::Tensor input) {
35  std::stringstream stream;
36  torch::save(input, stream);
37  torch::Tensor tensor;
38  torch::load(tensor, stream);
39  return tensor;
40 }
41 } // namespace
42 
43 TEST(SerializeTest, Basic) {
44  torch::manual_seed(0);
45 
46  auto x = torch::randn({5, 5});
47  auto y = save_and_load(x);
48 
49  ASSERT_TRUE(y.defined());
50  ASSERT_EQ(x.sizes().vec(), y.sizes().vec());
51  ASSERT_TRUE(x.allclose(y));
52 }
53 
54 TEST(SerializeTest, BasicToFile) {
55  torch::manual_seed(0);
56 
57  auto x = torch::randn({5, 5});
58 
59  auto tempfile = c10::make_tempfile();
60  torch::save(x, tempfile.name);
61 
62  torch::Tensor y;
63  torch::load(y, tempfile.name);
64 
65  ASSERT_TRUE(y.defined());
66  ASSERT_EQ(x.sizes().vec(), y.sizes().vec());
67  ASSERT_TRUE(x.allclose(y));
68 }
69 
70 TEST(SerializeTest, Resized) {
71  torch::manual_seed(0);
72 
73  auto x = torch::randn({11, 5});
74  x.resize_({5, 5});
75  auto y = save_and_load(x);
76 
77  ASSERT_TRUE(y.defined());
78  ASSERT_EQ(x.sizes().vec(), y.sizes().vec());
79  ASSERT_TRUE(x.allclose(y));
80 }
81 
82 TEST(SerializeTest, Sliced) {
83  torch::manual_seed(0);
84 
85  auto x = torch::randn({11, 5});
86  x = x.slice(0, 1, 5);
87  auto y = save_and_load(x);
88 
89  ASSERT_TRUE(y.defined());
90  ASSERT_EQ(x.sizes().vec(), y.sizes().vec());
91  ASSERT_TRUE(x.allclose(y));
92 }
93 
94 TEST(SerializeTest, NonContiguous) {
95  torch::manual_seed(0);
96 
97  auto x = torch::randn({11, 5});
98  x = x.slice(1, 1, 4);
99  auto y = save_and_load(x);
100 
101  ASSERT_TRUE(y.defined());
102  ASSERT_EQ(x.sizes().vec(), y.sizes().vec());
103  ASSERT_TRUE(x.allclose(y));
104 }
105 
106 TEST(SerializeTest, XOR) {
107  // We better be able to save and load an XOR model!
108  auto getLoss = [](Sequential model, uint32_t batch_size) {
109  auto inputs = torch::empty({batch_size, 2});
110  auto labels = torch::empty({batch_size});
111  for (size_t i = 0; i < batch_size; i++) {
112  inputs[i] = torch::randint(2, {2}, torch::kInt64);
113  labels[i] = inputs[i][0].item<int64_t>() ^ inputs[i][1].item<int64_t>();
114  }
115  auto x = model->forward<torch::Tensor>(inputs);
116  return torch::binary_cross_entropy(x, labels);
117  };
118 
119  auto model = xor_model();
120  auto model2 = xor_model();
121  auto model3 = xor_model();
122  auto optimizer = torch::optim::SGD(
123  model->parameters(),
124  torch::optim::SGDOptions(1e-1).momentum(0.9).nesterov(true).weight_decay(
125  1e-6));
126 
127  float running_loss = 1;
128  int epoch = 0;
129  while (running_loss > 0.1) {
130  torch::Tensor loss = getLoss(model, 4);
131  optimizer.zero_grad();
132  loss.backward();
133  optimizer.step();
134 
135  running_loss = running_loss * 0.99 + loss.sum().item<float>() * 0.01;
136  ASSERT_LT(epoch, 3000);
137  epoch++;
138  }
139 
140  auto tempfile = c10::make_tempfile();
141  torch::save(model, tempfile.name);
142  torch::load(model2, tempfile.name);
143 
144  auto loss = getLoss(model2, 100);
145  ASSERT_LT(loss.item<float>(), 0.1);
146 }
147 
148 TEST(SerializeTest, Optim) {
149  auto model1 = Linear(5, 2);
150  auto model2 = Linear(5, 2);
151  auto model3 = Linear(5, 2);
152 
153  // Models 1, 2, 3 will have the same parameters.
154  auto model_tempfile = c10::make_tempfile();
155  torch::save(model1, model_tempfile.name);
156  torch::load(model2, model_tempfile.name);
157  torch::load(model3, model_tempfile.name);
158 
159  auto param1 = model1->named_parameters();
160  auto param2 = model2->named_parameters();
161  auto param3 = model3->named_parameters();
162  for (const auto& p : param1) {
163  ASSERT_TRUE(p->allclose(param2[p.key()]));
164  ASSERT_TRUE(param2[p.key()].allclose(param3[p.key()]));
165  }
166 
167  // Make some optimizers with momentum (and thus state)
168  auto optim1 = torch::optim::SGD(
169  model1->parameters(), torch::optim::SGDOptions(1e-1).momentum(0.9));
170  auto optim2 = torch::optim::SGD(
171  model2->parameters(), torch::optim::SGDOptions(1e-1).momentum(0.9));
172  auto optim2_2 = torch::optim::SGD(
173  model2->parameters(), torch::optim::SGDOptions(1e-1).momentum(0.9));
174  auto optim3 = torch::optim::SGD(
175  model3->parameters(), torch::optim::SGDOptions(1e-1).momentum(0.9));
176  auto optim3_2 = torch::optim::SGD(
177  model3->parameters(), torch::optim::SGDOptions(1e-1).momentum(0.9));
178 
179  auto x = torch::ones({10, 5});
180 
181  auto step = [&x](torch::optim::Optimizer& optimizer, Linear model) {
182  optimizer.zero_grad();
183  auto y = model->forward(x).sum();
184  y.backward();
185  optimizer.step();
186  };
187 
188  // Do 2 steps of model1
189  step(optim1, model1);
190  step(optim1, model1);
191 
192  // Do 2 steps of model 2 without saving the optimizer
193  step(optim2, model2);
194  step(optim2_2, model2);
195 
196  // Do 2 steps of model 3 while saving the optimizer
197  step(optim3, model3);
198 
199  auto optim_tempfile = c10::make_tempfile();
200  torch::save(optim3, optim_tempfile.name);
201  torch::load(optim3_2, optim_tempfile.name);
202  step(optim3_2, model3);
203 
204  param1 = model1->named_parameters();
205  param2 = model2->named_parameters();
206  param3 = model3->named_parameters();
207  for (const auto& p : param1) {
208  const auto& name = p.key();
209  // Model 1 and 3 should be the same
210  ASSERT_TRUE(
211  param1[name].norm().item<float>() == param3[name].norm().item<float>());
212  ASSERT_TRUE(
213  param1[name].norm().item<float>() != param2[name].norm().item<float>());
214  }
215 }
216 
217 TEST(SerializeTest, XOR_CUDA) {
218  torch::manual_seed(0);
219  // We better be able to save and load a XOR model!
220  auto getLoss = [](Sequential model,
221  uint32_t batch_size,
222  bool is_cuda = false) {
223  auto inputs = torch::empty({batch_size, 2});
224  auto labels = torch::empty({batch_size});
225  if (is_cuda) {
226  inputs = inputs.cuda();
227  labels = labels.cuda();
228  }
229  for (size_t i = 0; i < batch_size; i++) {
230  inputs[i] = torch::randint(2, {2}, torch::kInt64);
231  labels[i] = inputs[i][0].item<int64_t>() ^ inputs[i][1].item<int64_t>();
232  }
233  auto x = model->forward<torch::Tensor>(inputs);
234  return torch::binary_cross_entropy(x, labels);
235  };
236 
237  auto model = xor_model();
238  auto model2 = xor_model();
239  auto model3 = xor_model();
240  auto optimizer = torch::optim::SGD(
241  model->parameters(),
242  torch::optim::SGDOptions(1e-1).momentum(0.9).nesterov(true).weight_decay(
243  1e-6));
244 
245  float running_loss = 1;
246  int epoch = 0;
247  while (running_loss > 0.1) {
248  torch::Tensor loss = getLoss(model, 4);
249  optimizer.zero_grad();
250  loss.backward();
251  optimizer.step();
252 
253  running_loss = running_loss * 0.99 + loss.sum().item<float>() * 0.01;
254  ASSERT_LT(epoch, 3000);
255  epoch++;
256  }
257 
258  auto tempfile = c10::make_tempfile();
259  torch::save(model, tempfile.name);
260  torch::load(model2, tempfile.name);
261 
262  auto loss = getLoss(model2, 100);
263  ASSERT_LT(loss.item<float>(), 0.1);
264 
265  model2->to(torch::kCUDA);
266  loss = getLoss(model2, 100, true);
267  ASSERT_LT(loss.item<float>(), 0.1);
268 
269  auto tempfile2 = c10::make_tempfile();
270  torch::save(model2, tempfile2.name);
271  torch::load(model3, tempfile2.name);
272 
273  loss = getLoss(model3, 100, true);
274  ASSERT_LT(loss.item<float>(), 0.1);
275 }
276 
277 TEST(
278  SerializeTest,
279  CanSerializeModulesWithIntermediateModulesWithoutParametersOrBuffers) {
280  struct C : torch::nn::Module {
281  C() {
282  register_buffer("foo", torch::ones(5, torch::kInt32));
283  }
284  };
285  struct B : torch::nn::Module {};
286  struct A : torch::nn::Module {
287  A() {
288  register_module("b", std::make_shared<B>());
289  register_module("c", std::make_shared<C>());
290  }
291  };
292  struct M : torch::nn::Module {
293  M() {
294  register_module("a", std::make_shared<A>());
295  }
296  };
297 
298  auto out = std::make_shared<M>();
299  std::stringstream ss;
300  torch::save(out, ss);
301  auto in = std::make_shared<M>();
302  torch::load(in, ss);
303 
304  const int output = in->named_buffers()["a.c.foo"].sum().item<int>();
305  ASSERT_EQ(output, 5);
306 }
void backward(c10::optional< Tensor > gradient=c10::nullopt, bool keep_graph=false, bool create_graph=false)
Computes the gradient of current tensor w.r.t. graph leaves.
Definition: TensorMethods.h:49
Definition: any.cpp:108
Optimizer that defines a required step() method that takes no arguments and produces no values...
Definition: optimizer.h:100
virtual void zero_grad()
Zeros out the gradients of all parameters.
Definition: optimizer.cpp:22
does bound shape inference given a C2 net.
The base class for all modules in PyTorch.
Definition: module.h:62
Definition: static.cpp:64
Definition: static.cpp:58
TempFile make_tempfile(std::string name_prefix="torch-file-")
Like try_make_tempfile, but throws an exception if a temporary file could not be returned.
Definition: tempfile.h:97