Skip to content
Snippets Groups Projects
Commit ff61f430 authored by Stephan Seitz's avatar Stephan Seitz
Browse files

Merge branch 'torch_native-rewrite'

parents 958e9427 78135130
Branches
Tags
No related merge requests found
Pipeline #18171 failed
...@@ -43,7 +43,7 @@ test_requires = ...@@ -43,7 +43,7 @@ test_requires =
pytest-html pytest-html
ansi2html ansi2html
pytest-cov pytest-cov
tensorflow tensorflow-gpu
# Require a specific Python version, e.g. Python 2.7 or >= 3.4 # 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.* # python_requires = >=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*
...@@ -62,7 +62,7 @@ testing = ...@@ -62,7 +62,7 @@ testing =
pytest-html pytest-html
ansi2html ansi2html
pytest-cov pytest-cov
tensorflow tensorflow-gpu
pybind11 = pybind11 =
cppimport cppimport
......
...@@ -35,7 +35,7 @@ Backward: ...@@ -35,7 +35,7 @@ Backward:
def __init__(self, def __init__(self,
forward_assignments: List[ps.Assignment], forward_assignments: List[ps.Assignment],
op_name: str = "AutoDiffOp", op_name: str = "autodiffop",
time_constant_fields: List[ps.Field] = [], time_constant_fields: List[ps.Field] = [],
constant_fields: List[ps.Field] = [], constant_fields: List[ps.Field] = [],
diff_fields_prefix='diff', # TODO: remove! diff_fields_prefix='diff', # TODO: remove!
...@@ -469,16 +469,17 @@ Backward: ...@@ -469,16 +469,17 @@ Backward:
return self.create_tensorflow_op(*args, backend='torch_native', **kwags) return self.create_tensorflow_op(*args, backend='torch_native', **kwags)
def create_tensorflow_op(self, def create_tensorflow_op(self,
inputfield_tensor_dict, inputfield_tensor_dict={},
forward_loop=None, forward_loop=None,
backward_loop=None, backward_loop=None,
use_cuda=True,
backend='tensorflow'): backend='tensorflow'):
""" """
Creates custom differentiable Tensorflow Op from assignments. 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 Will return either a single output tensor or a OrderedDict[field_name -> tf.Tensor] in case of multiple outputs
""" """
backend = backend.lower() 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) backend, AVAILABLE_BACKENDS)
additional_fields = [f for f in inputfield_tensor_dict.keys( additional_fields = [f for f in inputfield_tensor_dict.keys(
...@@ -536,17 +537,20 @@ Backward: ...@@ -536,17 +537,20 @@ Backward:
backward_loop = backward_function backward_loop = backward_function
if backend == 'tensorflow': if backend == 'tensorflow_native':
import pystencils_autodiff.backends._tensorflow import pystencils_autodiff.backends._tensorflow
op = pystencils_autodiff.backends._tensorflow.tensorflowop_from_autodiffop( op = pystencils_autodiff.backends._tensorflow.native_tensorflowop_from_autodiffop(self, use_cuda)
self, inputfield_tensor_dict, forward_loop, backward_loop)
elif backend == 'torch': elif backend == 'torch':
import pystencils_autodiff.backends._pytorch import pystencils_autodiff.backends._pytorch
op = pystencils_autodiff.backends._pytorch.create_autograd_function( op = pystencils_autodiff.backends._pytorch.create_autograd_function(
self, inputfield_tensor_dict, forward_loop, backward_loop) self, inputfield_tensor_dict, forward_loop, backward_loop)
elif backend == 'torch_native': elif backend == 'torch_native':
import pystencils_autodiff.backends._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: else:
raise NotImplementedError() raise NotImplementedError()
......
...@@ -6,4 +6,4 @@ a Torch or a Tensorflow operation or we can compile a static ...@@ -6,4 +6,4 @@ a Torch or a Tensorflow operation or we can compile a static
library to be directly loaded into Torch/Tensorflow. library to be directly loaded into Torch/Tensorflow.
""" """
AVAILABLE_BACKENDS = ['tensorflow', 'torch', 'tensorflow_cpp', 'torch_native'] AVAILABLE_BACKENDS = ['tensorflow', 'torch', 'tensorflow_native', 'torch_native']
from collections.abc import Iterable
import stringcase
import tensorflow as tf import tensorflow as tf
from tensorflow.compat.v1 import get_default_graph, py_func from tensorflow.compat.v1 import get_default_graph, py_func
import pystencils_autodiff import pystencils_autodiff
from pystencils_autodiff.backends.astnodes import TensorflowModule
from pystencils_autodiff.tensorflow_jit import _hash
_num_generated_ops = 0 _num_generated_ops = 0
...@@ -36,6 +41,44 @@ def _py_func(func, inp, Tout, stateful=False, name=None, grad=None): ...@@ -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) 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, def tensorflowop_from_autodiffop(autodiffop: pystencils_autodiff.AutoDiffOp,
inputfield_tensor_dict, inputfield_tensor_dict,
forward_function, forward_function,
......
"""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()
import uuid
from collections import OrderedDict from collections import OrderedDict
from pystencils_autodiff.backends._pytorch import numpy_dtype_to_torch from pystencils_autodiff.backends._pytorch import numpy_dtype_to_torch
from pystencils_autodiff.backends.astnodes import TorchModule from pystencils_autodiff.backends.astnodes import TorchModule
from pystencils_autodiff.tensorflow_jit import _hash
try: try:
import torch import torch
...@@ -10,48 +10,83 @@ except ImportError: ...@@ -10,48 +10,83 @@ except ImportError:
pass pass
def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict): def create_autograd_function(autodiff_obj, use_cuda):
field_to_tensor_dict = inputfield_to_tensor_dict field_to_tensor_dict = dict()
# Allocate output tensor for forward and backward pass # Allocate output tensor for forward and backward pass
for field in autodiff_obj.forward_output_fields + autodiff_obj.backward_output_fields: # for field in autodiff_obj.forward_output_fields + autodiff_obj.backward_output_fields:
field_to_tensor_dict[field] = torch.zeros( # field_to_tensor_dict[field] = torch.zeros(
*field.shape, # *field.shape,
dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype), # dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype),
device=list(inputfield_to_tensor_dict.values())[0].device) # 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)
if is_cuda: if use_cuda:
forward_ast = autodiff_obj.forward_ast_gpu forward_ast = autodiff_obj.forward_ast_gpu
backward_ast = autodiff_obj.backward_ast_gpu backward_ast = autodiff_obj.backward_ast_gpu
else: else:
forward_ast = autodiff_obj.forward_ast_cpu forward_ast = autodiff_obj.forward_ast_cpu
backward_ast = autodiff_obj.backward_ast_cpu backward_ast = autodiff_obj.backward_ast_cpu
op_name = autodiff_obj.op_name + str(uuid.uuid4()) op_name = f'{autodiff_obj.op_name}_{_hash(str(autodiff_obj).encode()).hexdigest()}'
compiled_op = TorchModule(op_name, [forward_ast, backward_ast]) module = TorchModule(op_name, [forward_ast, backward_ast])
compiled_op = module.compile()
output_tensors = OrderedDict({f.name: field_to_tensor_dict[f] for f in autodiff_obj.forward_output_fields})
backward_output_tensors = OrderedDict( # print(TorchModule(op_name, [forward_ast, backward_ast]))
{f.name: field_to_tensor_dict[f] for f in autodiff_obj.backward_output_fields}) # wrapper = module.atoms(WrapperFunction)
# forward_wrapper_ast = [w for w in wrapper if w.function_name == "call_" + forward_ast.function_name][0]
def forward(self, **input_tensors): # 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()]
self.save_for_backward(**input_tensors) # 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) 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): 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)} 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) 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 = type(op_name, (torch.autograd.Function,), {})
cls.forward = forward cls.forward = forward
......
...@@ -35,7 +35,7 @@ else: ...@@ -35,7 +35,7 @@ else:
_do_not_link_flag = '-c' _do_not_link_flag = '-c'
_output_flag = '-o' _output_flag = '-o'
_shared_object_flag = '/DLL' _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" _position_independent_flag = "/DTHIS_FLAG_DOES_NOTHING"
get_compiler_config()['command'] = 'cl.exe' get_compiler_config()['command'] = 'cl.exe'
config_env = get_compiler_config()['env'] if 'env' in get_compiler_config() else {} config_env = get_compiler_config()['env'] if 'env' in get_compiler_config() else {}
......
...@@ -13,7 +13,7 @@ import sympy ...@@ -13,7 +13,7 @@ import sympy
import pystencils import pystencils
from pystencils_autodiff import create_backward_assignments 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 from pystencils_autodiff.backends.astnodes import PybindModule, TorchModule
torch = pytest.importorskip('torch') torch = pytest.importorskip('torch')
......
import argparse import argparse
import itertools
import os import os
import numpy as np import numpy as np
...@@ -165,10 +166,10 @@ def test_tfmad_gradient_check_torch(): ...@@ -165,10 +166,10 @@ def test_tfmad_gradient_check_torch():
a, b, out = ps.fields("a, b, out: float[21,13]") a, b, out = ps.fields("a, b, out: float[21,13]")
cont = ps.fd.Diff(a, 0) - ps.fd.Diff(a, 1) - \ cont = 2 * ps.fd.Diff(a, 0) - 1.5 * ps.fd.Diff(a, 1) \
ps.fd.Diff(b, 0) + ps.fd.Diff(b, 1) - ps.fd.Diff(b, 0) + 3 * ps.fd.Diff(b, 1)
discretize = ps.fd.Discretization2ndOrder(dx=1) discretize = ps.fd.Discretization2ndOrder(dx=1)
discretization = discretize(cont) discretization = discretize(cont) + 1.2*a.center
assignment = ps.Assignment(out.center(), discretization) assignment = ps.Assignment(out.center(), discretization)
assignment_collection = ps.AssignmentCollection([assignment], []) assignment_collection = ps.AssignmentCollection([assignment], [])
...@@ -189,12 +190,109 @@ def test_tfmad_gradient_check_torch(): ...@@ -189,12 +190,109 @@ def test_tfmad_gradient_check_torch():
function = auto_diff.create_tensorflow_op({ function = auto_diff.create_tensorflow_op({
a: a_tensor, a: a_tensor,
b: b_tensor b: b_tensor
}, }, backend='torch')
backend='torch')
torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor]) 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): def get_curl(input_field: ps.Field, curl_field: ps.Field):
"""Return a ps.AssignmentCollection describing the calculation of """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) the curl given a 2d or 3d vector field [z,y,x](f) or [y,x](f)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment