Merge branch 'torch_native-rewrite'

setup.cfg

@@ -43,7 +43,7 @@ test_requires =
     pytest-html
     ansi2html
     pytest-cov
-    tensorflow
+    tensorflow-gpu
 # Require a specific Python version, e.g. Python 2.7 or >= 3.4
 # python_requires = >=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*
 
@@ -62,7 +62,7 @@ testing =
     pytest-html
     ansi2html
     pytest-cov
-    tensorflow
+    tensorflow-gpu
 pybind11 =
     cppimport
 

src/pystencils_autodiff/_autodiff.py

@@ -35,7 +35,7 @@ Backward:
 
     def __init__(self,
                  forward_assignments: List[ps.Assignment],
-                 op_name: str = "AutoDiffOp",
+                 op_name: str = "autodiffop",
                  time_constant_fields: List[ps.Field] = [],
                  constant_fields: List[ps.Field] = [],
                  diff_fields_prefix='diff',  # TODO: remove!
@@ -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()
 

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']

src/pystencils_autodiff/backends/_tensorflow.py

+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,44 @@ 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
+
+    autodiff_obj.op_name = f'{autodiff_obj.op_name}_hash{_hash(str(autodiff_obj).encode()).hexdigest()}'
+    if use_cuda:
+        autodiff_obj.op_name += '_cuda'
+    forward_ast.function_name = autodiff_obj.op_name + "_forward"
+    backward_ast.function_name = autodiff_obj.op_name + "_backward"
+    module = TensorflowModule(autodiff_obj.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})
+
+    try:
+        tf.RegisterGradient(stringcase.pascalcase("call_" + forward_ast.function_name))(
+            gradient_calculation
+        )
+    except Exception:
+        pass
+
+    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,

src/pystencils_autodiff/backends/_tensorflow_cpp.py

-"""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()

src/pystencils_autodiff/backends/_torch_native.py

-import uuid
 from collections import OrderedDict
 
 from pystencils_autodiff.backends._pytorch import numpy_dtype_to_torch
 from pystencils_autodiff.backends.astnodes import TorchModule
+from pystencils_autodiff.tensorflow_jit import _hash
 
 try:
     import torch
@@ -10,48 +10,83 @@ 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:
         forward_ast = autodiff_obj.forward_ast_cpu
         backward_ast = autodiff_obj.backward_ast_cpu
 
-    op_name = autodiff_obj.op_name + str(uuid.uuid4())
-    compiled_op = TorchModule(op_name, [forward_ast, backward_ast])
-
-    output_tensors = OrderedDict({f.name: field_to_tensor_dict[f] for f in autodiff_obj.forward_output_fields})
-    backward_output_tensors = OrderedDict(
-        {f.name: field_to_tensor_dict[f] for f in autodiff_obj.backward_output_fields})
-
-    def forward(self, **input_tensors):
-
-        self.save_for_backward(**input_tensors)
+    op_name = f'{autodiff_obj.op_name}_{_hash(str(autodiff_obj).encode()).hexdigest()}'
+    module = TorchModule(op_name, [forward_ast, backward_ast])
+    compiled_op = module.compile()
+
+    # print(TorchModule(op_name, [forward_ast, backward_ast]))
+    # wrapper = module.atoms(WrapperFunction)
+    # forward_wrapper_ast = [w for w in wrapper if w.function_name == "call_" + forward_ast.function_name][0]
+    # backward_wrapper_ast = [w for w in wrapper if w.function_name == "call_" + backward_ast.function_name][0]
+    # forward_parameters = [str(p.symbol) for p in forward_wrapper_ast.get_parameters()]
+    # 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})
+        assert all(f.shape == args[i].shape for i, f in enumerate(autodiff_obj.forward_input_fields))
+        assert all(f.strides == tuple(args[i].stride(j) for j in range(args[i].ndim))
+                   for i, f in enumerate(autodiff_obj.forward_input_fields))
+        for field in autodiff_obj.forward_output_fields:
+            field_to_tensor_dict[field] = torch.zeros(
+                field.shape,
+                dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype),
+                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)
 
         getattr(compiled_op, "call_" + forward_ast.function_name)(**input_tensors, **output_tensors)
 
-        return output_tensors.values()
+        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)}
-        saved = self.saved_tensors
-
+        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=grad_outputs[0].device)
+
+        backward_output_tensors = OrderedDict(
+            {f.name: field_to_tensor_dict[f] for f in autodiff_obj.backward_output_fields})
         getattr(compiled_op, "call_" + backward_ast.function_name)(**gradients, **saved, **backward_output_tensors)
 
-        return backward_output_tensors.values()
+        return tuple(backward_output_tensors.values())
 
     cls = type(op_name, (torch.autograd.Function,), {})
     cls.forward = forward

src/pystencils_autodiff/tensorflow_jit.py

@@ -35,7 +35,7 @@ else:
     _do_not_link_flag = '-c'
     _output_flag = '-o'
     _shared_object_flag = '/DLL'
-    _include_flags = ['/I' + sysconfig.get_paths()['include'], '/I' + get_pystencils_include_path()]
+    _include_flags = ['-I' + sysconfig.get_paths()['include'], '-I' + get_pystencils_include_path()]
     _position_independent_flag = "/DTHIS_FLAG_DOES_NOTHING"
     get_compiler_config()['command'] = 'cl.exe'
     config_env = get_compiler_config()['env'] if 'env' in get_compiler_config() else {}

tests/backends/test_torch_native_compilation.py

@@ -13,7 +13,7 @@ import sympy
 
 import pystencils
 from pystencils_autodiff import create_backward_assignments
-from pystencils_autodiff._file_io import write_cached_content, write_file
+from pystencils_autodiff._file_io import write_cached_content
 from pystencils_autodiff.backends.astnodes import PybindModule, TorchModule
 
 torch = pytest.importorskip('torch')

tests/test_tfmad.py

 import argparse
+import itertools
 import os
 
 import numpy as np
@@ -165,10 +166,10 @@ def test_tfmad_gradient_check_torch():
 
     a, b, out = ps.fields("a, b, out: float[21,13]")
 
-    cont = ps.fd.Diff(a, 0) - ps.fd.Diff(a, 1) - \
-        ps.fd.Diff(b, 0) + 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)
+    discretization = discretize(cont) + 1.2*a.center
 
     assignment = ps.Assignment(out.center(), discretization)
     assignment_collection = ps.AssignmentCollection([assignment], [])
@@ -189,12 +190,109 @@ def test_tfmad_gradient_check_torch():
     function = auto_diff.create_tensorflow_op({
         a: a_tensor,
         b: b_tensor
-    },
-                                              backend='torch')
+    }, backend='torch')
 
     torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])
 
 
+@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]")
+
+    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 = 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
+    }
+    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, gradient_check', itertools.product((False, True), repeat=3))
+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(), 1.2*a.center + 0.1*b[1, 0])
+    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')
+    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.compat.v1.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)