From 9398d6d740000fc1003e975869a3b0db97053095 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Thu, 19 Sep 2019 15:51:46 +0200
Subject: [PATCH] Make tensorflow_native kind of working

---
 src/pystencils_autodiff/_autodiff.py          | 16 ++--
 src/pystencils_autodiff/backends/__init__.py  |  2 +-
 .../backends/_tensorflow.py                   | 38 ++++++++++
 .../backends/_tensorflow_cpp.py               | 15 ----
 .../backends/_torch_native.py                 | 36 +++++----
 tests/test_tfmad.py                           | 74 +++++++++++++++++--
 6 files changed, 138 insertions(+), 43 deletions(-)
 delete mode 100644 src/pystencils_autodiff/backends/_tensorflow_cpp.py

diff --git a/src/pystencils_autodiff/_autodiff.py b/src/pystencils_autodiff/_autodiff.py
index b8f16fe..723749c 100644
--- a/src/pystencils_autodiff/_autodiff.py
+++ b/src/pystencils_autodiff/_autodiff.py
@@ -469,16 +469,17 @@ Backward:
         return self.create_tensorflow_op(*args, backend='torch_native', **kwags)
 
     def create_tensorflow_op(self,
-                             inputfield_tensor_dict,
+                             inputfield_tensor_dict={},
                              forward_loop=None,
                              backward_loop=None,
+                             use_cuda=True,
                              backend='tensorflow'):
         """
         Creates custom differentiable Tensorflow Op from assignments.
         Will return either a single output tensor or a OrderedDict[field_name -> tf.Tensor] in case of multiple outputs
         """
         backend = backend.lower()
-        assert backend in AVAILABLE_BACKENDS, "\"{}\" is not a. Available backends: {}".format(
+        assert backend in AVAILABLE_BACKENDS, "\"{}\" is not a valid backend. Available backends: {}".format(
             backend, AVAILABLE_BACKENDS)
 
         additional_fields = [f for f in inputfield_tensor_dict.keys(
@@ -536,17 +537,20 @@ Backward:
 
             backward_loop = backward_function
 
-        if backend == 'tensorflow':
+        if backend == 'tensorflow_native':
             import pystencils_autodiff.backends._tensorflow
-            op = pystencils_autodiff.backends._tensorflow.tensorflowop_from_autodiffop(
-                self, inputfield_tensor_dict, forward_loop, backward_loop)
+            op = pystencils_autodiff.backends._tensorflow.native_tensorflowop_from_autodiffop(self, use_cuda)
         elif backend == 'torch':
             import pystencils_autodiff.backends._pytorch
             op = pystencils_autodiff.backends._pytorch.create_autograd_function(
                 self, inputfield_tensor_dict, forward_loop, backward_loop)
         elif backend == 'torch_native':
             import pystencils_autodiff.backends._torch_native
-            op = pystencils_autodiff.backends._torch_native.create_autograd_function(self, inputfield_tensor_dict)
+            op = pystencils_autodiff.backends._torch_native.create_autograd_function(self, use_cuda)
+        elif backend == 'tensorflow':
+            import pystencils_autodiff.backends._tensorflow
+            op = pystencils_autodiff.backends._tensorflow.tensorflowop_from_autodiffop(
+                self, inputfield_tensor_dict, forward_loop, backward_loop)
         else:
             raise NotImplementedError()
 
diff --git a/src/pystencils_autodiff/backends/__init__.py b/src/pystencils_autodiff/backends/__init__.py
index d22da0f..8744d19 100644
--- a/src/pystencils_autodiff/backends/__init__.py
+++ b/src/pystencils_autodiff/backends/__init__.py
@@ -6,4 +6,4 @@ a Torch or a Tensorflow operation or we can compile a static
 library to be directly loaded into Torch/Tensorflow.
 """
 
-AVAILABLE_BACKENDS = ['tensorflow', 'torch', 'tensorflow_cpp', 'torch_native']
+AVAILABLE_BACKENDS = ['tensorflow', 'torch', 'tensorflow_native', 'torch_native']
diff --git a/src/pystencils_autodiff/backends/_tensorflow.py b/src/pystencils_autodiff/backends/_tensorflow.py
index b243437..a5ca03b 100644
--- a/src/pystencils_autodiff/backends/_tensorflow.py
+++ b/src/pystencils_autodiff/backends/_tensorflow.py
@@ -1,7 +1,12 @@
+from collections.abc import Iterable
+
+import stringcase
 import tensorflow as tf
 from tensorflow.compat.v1 import get_default_graph, py_func
 
 import pystencils_autodiff
+from pystencils_autodiff.backends.astnodes import TensorflowModule
+from pystencils_autodiff.tensorflow_jit import _hash
 
 _num_generated_ops = 0
 
@@ -36,6 +41,39 @@ def _py_func(func, inp, Tout, stateful=False, name=None, grad=None):
         return py_func(func, inp, Tout, stateful=stateful, name=name)
 
 
+def native_tensorflowop_from_autodiffop(autodiff_obj: pystencils_autodiff.AutoDiffOp,
+                                        use_cuda):
+
+    if use_cuda:
+        forward_ast = autodiff_obj.forward_ast_gpu
+        backward_ast = autodiff_obj.backward_ast_gpu
+    else:
+        forward_ast = autodiff_obj.forward_ast_cpu
+        backward_ast = autodiff_obj.backward_ast_cpu
+
+    op_name = f'{autodiff_obj.op_name}_{_hash(str(autodiff_obj).encode()).hexdigest()}'
+    forward_ast.function_name = autodiff_obj.op_name + "_forward"
+    backward_ast.function_name = autodiff_obj.op_name + "_backward"
+    module = TensorflowModule(op_name, [forward_ast, backward_ast])
+    compiled_op = module.compile()
+
+    backward_func = getattr(compiled_op, stringcase.snakecase(
+        stringcase.pascalcase("call_" + backward_ast.function_name)))
+
+    def gradient_calculation(op, grad):
+        if isinstance(grad, Iterable):
+            grad = [grad]
+        return backward_func(**{autodiff_obj.backward_input_fields[i].name: g for i, g in enumerate(grad)},
+                             **{autodiff_obj.forward_input_fields[i].name: inp for i, inp in enumerate(op.inputs)
+                                if autodiff_obj.forward_input_fields[i] in backward_ast.fields_accessed})
+
+    tf.RegisterGradient(stringcase.pascalcase("call_" + forward_ast.function_name))(
+        gradient_calculation
+    )
+
+    return getattr(compiled_op, stringcase.snakecase(stringcase.pascalcase("call_" + forward_ast.function_name)))
+
+
 def tensorflowop_from_autodiffop(autodiffop: pystencils_autodiff.AutoDiffOp,
                                  inputfield_tensor_dict,
                                  forward_function,
diff --git a/src/pystencils_autodiff/backends/_tensorflow_cpp.py b/src/pystencils_autodiff/backends/_tensorflow_cpp.py
deleted file mode 100644
index d64f176..0000000
--- a/src/pystencils_autodiff/backends/_tensorflow_cpp.py
+++ /dev/null
@@ -1,15 +0,0 @@
-"""Implementing a custom Tensorflow Op in C++ has some advantages and disadvantages
-
-Advantages:
-- GPU support without any hacks
-- Access to raw tensors without conversion to numpy
-- Custom Ops will be serializable
-
-Disadavantages:
-- C++ Code has to be build with correct parameters and ABI
-for present Tensorflow version (best integrated into Tensorflow build)
-
-"""
-
-
-# raise NotImplementedError()
diff --git a/src/pystencils_autodiff/backends/_torch_native.py b/src/pystencils_autodiff/backends/_torch_native.py
index b574c30..b8b2ff5 100644
--- a/src/pystencils_autodiff/backends/_torch_native.py
+++ b/src/pystencils_autodiff/backends/_torch_native.py
@@ -10,20 +10,16 @@ except ImportError:
     pass
 
 
-def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict):
-    field_to_tensor_dict = inputfield_to_tensor_dict
-
+def create_autograd_function(autodiff_obj, use_cuda):
+    field_to_tensor_dict = dict()
     # Allocate output tensor for forward and backward pass
-    for field in autodiff_obj.forward_output_fields + autodiff_obj.backward_output_fields:
-        field_to_tensor_dict[field] = torch.zeros(
-            *field.shape,
-            dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype),
-            device=list(inputfield_to_tensor_dict.values())[0].device)
-
-    all_tensors = field_to_tensor_dict.values()
-    is_cuda = all(a.is_cuda for a in all_tensors)
+    # for field in autodiff_obj.forward_output_fields + autodiff_obj.backward_output_fields:
+    # field_to_tensor_dict[field] = torch.zeros(
+    # *field.shape,
+    # dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype),
+    # device=list(inputfield_to_tensor_dict.values())[0].device)
 
-    if is_cuda:
+    if use_cuda:
         forward_ast = autodiff_obj.forward_ast_gpu
         backward_ast = autodiff_obj.backward_ast_gpu
     else:
@@ -42,6 +38,12 @@ def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict):
     # backward_parameters = [str(p.symbol) for p in backward_wrapper_ast.get_parameters()]
 
     def forward(self, *args):
+
+        if use_cuda:
+            args = [a.contiguous().cuda() for a in args]
+        else:
+            args = [a.contiguous().cpu() for a in args]
+
         input_tensors = dict()
         input_tensors.update({f.name: args[i] for i, f in enumerate(
             autodiff_obj.forward_input_fields) if f in forward_ast.fields_accessed})
@@ -52,7 +54,7 @@ def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict):
             field_to_tensor_dict[field] = torch.zeros(
                 field.shape,
                 dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype),
-                device=list(inputfield_to_tensor_dict.values())[0].device)
+                device=args[0].device)
         output_tensors = OrderedDict({f.name: field_to_tensor_dict[f] for f in autodiff_obj.forward_output_fields})
 
         self.save_for_backward(*args)
@@ -62,17 +64,23 @@ def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict):
         return tuple(output_tensors.values())
 
     def backward(self, *grad_outputs):
+        if use_cuda:
+            grad_outputs = [a.contiguous().cuda() for a in grad_outputs]
+        else:
+            grad_outputs = [a.contiguous().cpu() for a in grad_outputs]
         gradients = {f.name: grad_outputs[i] for i, f in enumerate(autodiff_obj.backward_input_fields)}
         assert all(f.shape == grad_outputs[i].shape for i, f in enumerate(autodiff_obj.backward_input_fields))
         assert all(f.strides == tuple(grad_outputs[i].stride(j) for j in range(grad_outputs[i].ndim))
                    for i, f in enumerate(autodiff_obj.backward_input_fields))
+        assert all(a.is_cuda == use_cuda for a in grad_outputs), f"Some of the tensors where on the wrong device. " \
+            f"Op was compiled for CUDA: {str(use_cuda)}"
         saved = {f.name: self.saved_tensors[i] for i, f in enumerate(
             autodiff_obj.forward_input_fields) if f in backward_ast.fields_accessed}
         for field in autodiff_obj.backward_output_fields:
             field_to_tensor_dict[field] = torch.zeros(
                 field.shape,
                 dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype),
-                device=list(inputfield_to_tensor_dict.values())[0].device)
+                device=grad_outputs[0].device)
 
         backward_output_tensors = OrderedDict(
             {f.name: field_to_tensor_dict[f] for f in autodiff_obj.backward_output_fields})
diff --git a/tests/test_tfmad.py b/tests/test_tfmad.py
index 7065d9c..46279dd 100644
--- a/tests/test_tfmad.py
+++ b/tests/test_tfmad.py
@@ -1,4 +1,5 @@
 import argparse
+import itertools
 import os
 
 import numpy as np
@@ -165,8 +166,8 @@ def test_tfmad_gradient_check_torch():
 
     a, b, out = ps.fields("a, b, out: float[21,13]")
 
-    cont = 2*ps.fd.Diff(a, 0) - 1.5*ps.fd.Diff(a, 1) - \
-        ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1)
+    cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) \
+        - ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1)
     discretize = ps.fd.Discretization2ndOrder(dx=1)
     discretization = discretize(cont) + 1.2*a.center
 
@@ -194,9 +195,10 @@ def test_tfmad_gradient_check_torch():
     torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])
 
 
-@pytest.mark.parametrize('with_offsets', (True, False))
-def test_tfmad_gradient_check_torch_native(with_offsets):
+@pytest.mark.parametrize('with_offsets, with_cuda', itertools.product((False, True), repeat=2))
+def test_tfmad_gradient_check_torch_native(with_offsets, with_cuda):
     torch = pytest.importorskip('torch')
+    import torch
 
     a, b, out = ps.fields("a, b, out: float64[21,13]")
 
@@ -223,17 +225,75 @@ def test_tfmad_gradient_check_torch_native(with_offsets):
     a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True).contiguous()
     b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True).contiguous()
 
+    if with_cuda:
+        a_tensor = a_tensor.cuda()
+        b_tensor = b_tensor.cuda()
+
+    function = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='torch_native')
+
     dict = {
         a: a_tensor,
         b: b_tensor
     }
-    function = auto_diff.create_tensorflow_op(dict, backend='torch_native')
-
-    import torch
     torch.autograd.gradcheck(function.apply, tuple(
         [dict[f] for f in auto_diff.forward_input_fields]), atol=1e-4, raise_exception=True)
 
 
+# @pytest.mark.parametrize('with_offsets, with_cuda', itertools.product((False, True), repeat=2))
+@pytest.mark.parametrize('with_offsets, with_cuda, gradient_check', ((False, False, True),))
+def test_tfmad_gradient_check_tensorflow_native(with_offsets, with_cuda, gradient_check):
+    pytest.importorskip('tensorflow')
+    import tensorflow as tf
+
+    a, b, out = ps.fields("a, b, out: double[21,13]")
+    print(a.shape)
+
+    if with_offsets:
+        cont = 2*ps.fd.Diff(a, 0) - 1.5*ps.fd.Diff(a, 1) - ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1)
+        discretize = ps.fd.Discretization2ndOrder(dx=1)
+        discretization = discretize(cont)
+
+        assignment = ps.Assignment(out.center(), discretization + 1.2*a.center())
+    else:
+        assignment = ps.Assignment(out.center(), 1.2*a.center + 0.1*b.center)
+
+    assignment_collection = ps.AssignmentCollection([assignment], [])
+    print('Forward')
+    print(assignment_collection)
+
+    print('Backward')
+    auto_diff = pystencils_autodiff.AutoDiffOp(assignment_collection,
+                                               diff_mode='transposed-forward')
+    backward = auto_diff.backward_assignments
+    print(backward)
+    print('Forward output fields (to check order)')
+    print(auto_diff.forward_input_fields)
+
+    a_tensor = tf.Variable(np.zeros(a.shape, a.dtype.numpy_dtype))
+    b_tensor = tf.Variable(np.zeros(a.shape, a.dtype.numpy_dtype))
+    # out_tensor = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='tensorflow_native')
+    # print(out_tensor)
+
+    out_tensor = auto_diff.create_tensorflow_op(use_cuda=with_cuda, backend='tensorflow_native')(a=a_tensor, b=b_tensor)
+
+    with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        sess.run(out_tensor)
+
+        if gradient_check:
+            gradient_error = compute_gradient_error_without_border(
+                [a_tensor, b_tensor], [a.shape, b.shape],
+                out_tensor,
+                out.shape,
+                num_border_pixels=2,
+                ndim=2,
+                debug=False)
+            print('error: %s' % gradient_error.max_error)
+            print('avg error: %s' % gradient_error.avg_error)
+
+            assert any(e < 1e-4 for e in gradient_error.values())
+
+
 def get_curl(input_field: ps.Field, curl_field: ps.Field):
     """Return a ps.AssignmentCollection describing the calculation of
     the curl given a 2d or 3d vector field [z,y,x](f) or [y,x](f)
-- 
GitLab