Added chainer.grad features to chainerx.grad

dido1998 · dido1998 · Jul 23, 2019 · Jun 24, 2019 · Jul 2, 2019 · Jul 2, 2019
commit dc1dbf3e98ea2ec26cc9cd6c3f3ffef463078385
diff --git a/chainer/function_node.py b/chainer/function_node.py
@@ -1011,39  1011,36 @@ def grad(outputs, inputs, grad_outputs=None, grad_inputs=None, set_grad=False,
 
     # Check if all the inputs are chainerx arrays and if so
     # Relies in chainerx.grad function
-    chx_inputs = sum(map(lambda x: x._has_chainerx_array
-                         if x is not None else False, inputs))
-    if chx_inputs == len(inputs):
-        # Need to access the arrays to invoke the chainer grad function
-        if grad_inputs is not None:
-            grad_inputs_chx = list(map(lambda x: x._data[0], grad_inputs))
-        else:
-            grad_inputs_chx = []
-        if grad_outputs is not None:
-            grad_outs_chx = list(map(lambda x: x._data[0], grad_outputs))
-        else:
-            grad_outs_chx = []
-        if retain_grad:
-            raise ValueError(
-                'retain_grad is not supported on chainerx.grad interface')
     n_chx_inputs = sum([False if x is None else x._has_chainerx_array
                         for x in inputs])
     if n_chx_inputs == len(inputs):
         if loss_scale is not None:
             raise ValueError(
                 'loss_scale is not supported on chainerx.grad interface')
 
-        outputs_chx = list(map(lambda x: x._data[0], outputs))
-        inputs_chx = list(map(lambda x: x._data[0], inputs))
         # Need to access the arrays to invoke the chainer grad function
         if grad_inputs:
             grad_inputs_chx = [x._data[0] for x in grad_inputs]
         else:
             grad_inputs_chx = []
         if grad_outputs:
             grad_outputs_chx = [x._data[0] for x in grad_outputs]
         else:
             grad_outputs_chx = []
         outputs_chx = [x._data[0] for x in outputs]
         inputs_chx = [x._data[0] for x in inputs]
         grads = chainerx.grad(outputs_chx, inputs_chx,
                               backprop_id=None,
                               enable_double_backprop=enable_double_backprop,
                               set_grad=set_grad,
                               retain_grad=retain_grad,
                               grad_inputs=grad_inputs_chx,
-                              grad_outputs=grad_outs_chx)
-        ret_vars = [variable.Variable(g, requires_grad=False) for g in grads]
-        if set_grad:
-            for x, g in zip(inputs, ret_vars):
-                x.grad_var = g
                               grad_outputs=grad_outputs_chx)
         ret_vars = [variable.Variable(g,
                                       requires_grad=g.is_backprop_required())
                     for g in grads]
         return ret_vars
-    elif chx_inputs > 0:
     elif n_chx_inputs > 0:
         raise TypeError(
             'Mixing chainerx and non-chainerx variables is not allowed')
 

diff --git a/chainerx_cc/chainerx/backward.cc b/chainerx_cc/chainerx/backward.cc
@@ -207,12  207,14 @@ class BackwardImpl {
             const std::vector<ConstArrayRef>& outputs,
             const BackpropId& backprop_id,
             DoubleBackpropOption double_backprop,
-            std::unordered_map<ArrayNode*, internal::GradRef> array_node_grad_map)
             std::unordered_map<ArrayNode*, internal::GradRef> array_node_grad_map,
             bool retain_grad = false)
         : inputs_{inputs},
           outputs_{outputs},
           backprop_id_{backprop_id},
           double_backprop_{double_backprop},
-          array_node_grad_map_{std::move(array_node_grad_map)} {
           array_node_grad_map_{std::move(array_node_grad_map)},
           retain_grad_{retain_grad} {
         if (!outputs_.empty()) {
             // Collect output array nodes (can be nullptr).
             output_array_nodes_.reserve(outputs.size());
@@ -396,7  398,11 @@ class BackwardImpl {
                 if (output_array_node != nullptr) {
                     std::shared_ptr<ArrayBody> body = output_array_node->weak_body().lock();
                     if (body != nullptr && !body->IsGradRequired(backprop_id_)) {
-                        body->ClearGrad(backprop_id_);
                         if (retain_grad_) {
                             internal::ArrayBody::RequireGrad(body, backprop_id_);
                         } else {
                             body->ClearGrad(backprop_id_);
                         }
                     }
                 }
             }
@@ -567,6  573,9 @@ class BackwardImpl {
     // Represents the subgraph required for backprop in case any inputs are specified.
     // Specifically, stores for an op nodes a vector of boolean flags whether an input at that index is included in the subgraph.
     std::unordered_map<OpNode*, std::vector<uint8_t>> input_required_flags_;
 
     // To mark if grads of intermediate nodes in the graph should be kept
     bool retain_grad_;
 };
 
 }  // namespace
@@ -586,56  595,59 @@ void Backward(
     BackwardImpl{{}, outputs, actual_backprop_id, double_backprop}.Run();
 }
 
-
 std::vector<nonstd::optional<Array>> Grad(
         const std::vector<ConstArrayRef>& outputs,
         const std::vector<ConstArrayRef>& inputs,
         const nonstd::optional<BackpropId>& backprop_id,
         DoubleBackpropOption double_backprop,
         bool set_grad,
         bool retain_grad,
         const std::vector<ConstArrayRef>& grad_inputs,
-        const std::vector<ConstArrayRef>& grad_outputs){
-
         const std::vector<ConstArrayRef>& grad_outputs) {
     if (inputs.empty()) {
         return {};
     }
     if (outputs.empty()) {
         return std::vector<nonstd::optional<Array>>(inputs.size(), nonstd::nullopt);
     }
 
-    CHAINERX_ASSERT(!grad_inputs.size() || (inputs.size() == grad_inputs.size()));
-    CHAINERX_ASSERT(!grad_outputs.size() || (outputs.size() == grad_outputs.size()));
     if (!grad_inputs.empty() && inputs.size() != grad_inputs.size()) {
         throw GradientError{"Grad inputs and inputs must have the same size"};
     }
     if (!grad_outputs.empty() && outputs.size() != grad_outputs.size()) {
         throw GradientError{"Grad_outputs and outputs must have the same size"};
     }
 
     std::vector<nonstd::optional<Array>> input_grads;
     std::vector<nonstd::optional<Array>> output_grads;
     input_grads.reserve(inputs.size());
     output_grads.reserve(outputs.size());
 
     BackpropId actual_backprop_id = internal::GetArrayBackpropId(outputs.front().get(), backprop_id);
-    
 
     // Initialize the grad map with newly created gradient arrays of the inputs.
     // The existing gradients of the inputs are thus not modified.
     std::unordered_map<ArrayNode*, internal::GradRef> array_node_grad_map;
-    for (size_t i=0; i<inputs.size();i  ) {
     for (size_t i = 0; i < inputs.size(); i  ) {
         const std::shared_ptr<ArrayBody>& array_body = internal::GetArrayBody(inputs[i]);
         if (const std::shared_ptr<ArrayNode>& input_array_node = array_body->GetArrayNode(actual_backprop_id)) {
-            input_grads.emplace_back(grad_inputs.size()? nonstd::optional<Array>{grad_inputs[i]} : nonstd::optional<Array>{});
             input_grads.emplace_back(!grad_inputs.empty() ? nonstd::optional<Array>{grad_inputs[i]} : nonstd::optional<Array>{});
             array_node_grad_map.emplace(input_array_node.get(), internal::GradRef{&input_grads.back()});
         } else {
             input_grads.emplace_back(nonstd::nullopt);
         }
     }
 
-    // Push initial output grads, Run assigns to 1 in other case 
-    for (size_t i=0; i<grad_outputs.size();i  ) {
     // Push initial output grads, Run assigns to 1 in other case
     for (size_t i = 0; i < grad_outputs.size(); i  ) {
         const std::shared_ptr<ArrayBody>& array_body = internal::GetArrayBody(outputs[i]);
         if (const std::shared_ptr<ArrayNode>& array_node = array_body->GetArrayNode(actual_backprop_id)) {
             output_grads.emplace_back(grad_outputs[i]);
             array_node_grad_map.emplace(array_node.get(), internal::GradRef{&output_grads.back()});
         }
     }
 
-    BackwardImpl{inputs, outputs, actual_backprop_id, double_backprop, std::move(array_node_grad_map)}.Run();
     BackwardImpl{inputs, outputs, actual_backprop_id, double_backprop, std::move(array_node_grad_map), retain_grad}.Run();
 
     size_t i = 0;
     // input_grads may contain unmodified array bodies (nullptr) for arrays that are not included in the graph.

diff --git a/chainerx_cc/chainerx/backward.h b/chainerx_cc/chainerx/backward.h
@@ -51,6  51,7 @@ std::vector<nonstd::optional<Array>> Grad(
         const nonstd::optional<BackpropId>& backprop_id = nonstd::nullopt,
         DoubleBackpropOption double_backprop = DoubleBackpropOption::kDisable,
         bool set_grad = false,
         bool retain_grad = false,
         const std::vector<ConstArrayRef>& grad_inputs = {},
         const std::vector<ConstArrayRef>& grad_outputs = {});
 

diff --git a/chainerx_cc/chainerx/python/backward.cc b/chainerx_cc/chainerx/python/backward.cc
@@ -64,8  64,9 @@ void InitChainerxBackward(pybind11::module& m) {
              const nonstd::optional<BackpropId>& backprop_id,
              bool enable_double_backprop,
              bool set_grad,
              bool retain_grad,
              const std::vector<ArrayBodyPtr>& grad_inputs,
-             const std::vector<ArrayBodyPtr>& grad_outputs){
              const std::vector<ArrayBodyPtr>& grad_outputs) {
               std::vector<Array> output_arrays = ConvertToArrays(outputs);
               std::vector<Array> input_arrays = ConvertToArrays(inputs);
 
@@ -79,6  80,7 @@ void InitChainerxBackward(pybind11::module& m) {
                            backprop_id,
                            double_backprop,
                            set_grad,
                            retain_grad,
                            std::vector<ConstArrayRef>{grad_input_arrays.begin(), grad_input_arrays.end()},
                            std::vector<ConstArrayRef>{grad_output_arrays.begin(), grad_output_arrays.end()});
               return internal::MoveArrayBodies(std::move(grads));
@@ -88,6  90,7 @@ void InitChainerxBackward(pybind11::module& m) {
           "backprop_id"_a = nullptr,
           "enable_double_backprop"_a = false,
           "set_grad"_a = false,
           "retain_grad"_a = false,
           "grad_inputs"_a = std::vector<ArrayBodyPtr>{},
           "grad_outputs"_a = std::vector<ArrayBodyPtr>{});
 }

diff --git a/tests/chainer_tests/test_function_node.py b/tests/chainer_tests/test_function_node.py
@@ -883,8  883,6 @@ def _init_attrs(self, names):
         for name in names:
             v = chainer.Variable(
                 numpy.random.randint(-4, 6, self.shape).astype('f'), name=name)
-            if self.extend_graph_x:
-                v *= 1.
             ret.append(v)
             setattr(self, name, v)
         return ret
@@ -920,6  918,11 @@ def use_gpu(self):
             if isinstance(value, chainer.Variable):
                 value.to_gpu()
 
     def use_chainerx(self, device='native:1'):
         for value in six.itervalues(self.__dict__):
             if isinstance(value, chainer.Variable):
                 value.to_device(backend.get_device(device))
 
     def forward(self):
         raise NotImplementedError
 
@@ -943,6  946,13 @@ def check_grad(self):
         ys = [getattr(self, name) for name in self.y_names]
         if self.extend_graph_y:
             self._ys = [v * 1. for v in ys]
 
         # graph_x extension should be done here
         # to avoid chainer/chainerx mixed graph
         if self.extend_graph_x:
             for v in self.xs:
                 v *= 1.
 
         gxs = chainer.grad(ys, self.xs, self.gys, self.gxs,
                            loss_scale=self.loss_scale)
 
@@ -960,14  970,29 @@ def check_grad(self):
             self._print_variables('gxs (expected)', expected)
             raise
 
     def chainerx_grad(self, device):
         if self.loss_scale:
             pytest.skip('Unsupported configurations')
         self.use_chainerx(device)
         self.check_grad()
 
     def test_grad_cpu(self):
         self.check_grad()
 
     @attr.chainerx
     def test_grad_chainerx_cpu(self):
         self.chainerx_grad('native:1')
 
     @attr.gpu
     def test_grad_gpu(self):
         self.use_gpu()
         self.check_grad()
 
     @attr.chainerx
     @attr.gpu
     def test_grad_chainerx_gpu(self):
         self.chainerx_grad('cuda:0')
 
     def check_double_grad(self):
         self.forward()
         ys = [getattr(self, name) for name in self.y_names]
@@ -992,11  1017,26 @@ def check_double_grad(self):
     def test_double_grad_cpu(self):
         self.check_double_grad()
 
     def double_chainerx_grad(self, device):
         if self.loss_scale:
             pytest.skip('Unsupported configurations')
         self.use_chainerx(device)
         self.check_double_grad()
 
     @attr.chainerx
     def test_double_grad_chainerx_cpu(self):
         self.double_chainerx_grad('native:1')
 
     @attr.gpu
     def test_double_grad_gpu(self):
         self.use_gpu()
         self.check_double_grad()
 
     @attr.chainerx
     @attr.gpu
     def test_double_grad_chainerx_gpu(self):
         self.double_chainerx_grad('cuda:0')
 
 
 @testing.parameterize(*testing.product({
     'loss_scale': [None, 1, 10],