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