Let script module buffer attributes can also cast device/type (#19700)

wanchaol · facebook-github-bot · commit 236c2b238724 · 2019-04-26T13:06:52.000-07:00
Summary: Tested locally this fix #19039, did not add a test since there's no way to create a script module in the cpp world. Pull Request resolved: #19700 Differential Revision: D15094195 Pulled By: wanchaol fbshipit-source-id: fcc2c1e5efbc160d976ae485ba2457442f62f065
diff --git a/test/cpp/jit/test.cpp b/test/cpp/jit/test.cpp
@@ -80,6 +80,7 @@ namespace jit {
   _(ArgumentSpec)               \
   _(Fusion)                     \
   _(GraphExecutor)              \
+  _(ModuleConversion)           \
   _(Interp)
 
 #if defined(USE_GTEST)
diff --git a/test/cpp/jit/test_misc.h b/test/cpp/jit/test_misc.h
@@ -758,6 +758,30 @@ void testModuleDefine() {
   AT_ASSERT(result.toTensor().item<float>() == 6)
 }
 
+void testModuleConversion() {
+  auto m = std::make_shared<script::Module>();
+  {
+    // test cuda to cpu for params and buffers
+    m->register_parameter("foo", torch::ones({}, at::kCUDA), false);
+    m->register_buffer("bar", torch::ones({}, at::kCUDA));
+    
+    m->to(at::kCUDA);
+    m->to(at::kCPU);
+    AT_ASSERT(m->get_parameter("foo").data().device().is_cpu());
+    AT_ASSERT(m->get_buffer("bar").data().device().is_cpu());
+  }
+  {
+    // test cpu to cuda for params and buffers
+    m->register_parameter("foo", torch::ones({}), false);
+    m->register_buffer("bar", torch::ones({}));
+    
+    m->to(at::kCUDA);
+    AT_ASSERT(m->get_parameter("foo").data().device().is_cuda());
+    AT_ASSERT(m->get_buffer("bar").data().device().is_cuda());
+  }
+}
+
+
 static int testPassValue = 0;
 void fakePass(std::shared_ptr<Graph>& g) {
   testPassValue++;
diff --git a/torch/csrc/jit/script/module.cpp b/torch/csrc/jit/script/module.cpp
@@ -101,25 +101,39 @@ void Module::save(
   ExportModule(*this, filename, extra_files);
 }
 
-void Module::to_impl(
+void module_state_to(
+    const Slot& s,
     const c10::optional<at::Device>& device,
     const c10::optional<at::ScalarType>& dtype,
     bool non_blocking) {
-  // First call `to()` on every child module.
-  for (auto& child : get_modules()) {
-    child->to_impl(device, dtype, non_blocking);
-  }
-  // Then convert every of our parameters.
-  for (auto& parameter : get_parameters()) {
     // Need to access the `at::Tensor` as a `Variable` here.
-    autograd::Variable variable = parameter.value().toTensor();
+    autograd::Variable variable = s.value().toTensor();
     at::Tensor data = variable.data();
     // Use the data's original device or dtype if not supplied here.
     auto new_data = data.to(
         device.value_or(data.device()),
         dtype.value_or(data.scalar_type()),
         non_blocking);
     variable.set_data(new_data);
+}
+
+void Module::to_impl(
+    const c10::optional<at::Device>& device,
+    const c10::optional<at::ScalarType>& dtype,
+    bool non_blocking) {
+  // First call `to()` on every child module.
+  for (auto& child : get_modules()) {
+    child->to_impl(device, dtype, non_blocking);
+  }
+  // Then convert every of our parameters.
+  for (auto& parameter : get_parameters()) {
+    module_state_to(parameter, device, dtype, non_blocking);
+  }
+  // Then convert every tensor attributes (buffers).
+  for (auto& attr : get_attributes()) {
+    if (attr.type()->isSubtypeOf(TensorType::get())) {
+      module_state_to(attr, device, dtype, non_blocking);
+    }
   }
 }
 
