Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support RayTransform custom backends #1540

Merged
merged 39 commits into from
Apr 16, 2020

Conversation

adriaangraas
Copy link
Contributor

At this point it is difficult from the user level to specify a different backend for the RayTransform operator. (Current backends are Scikit-image, ASTRA CPU and ASTRA CUDA). Loading a custom backend is necessary if you like to modify the way ASTRA is called, for instance.

In this PR I propose to move some backend logic out of the RayTransformBase and RayTransform classes and into new RayTransformImplBase class and implementation classes. Something like this was already in place for impl='astra_cuda'. With this PR a user can for instance subclass the ASTRA CPU implementation and feed its type into the impl argument (strategy design pattern).

    class MyRayTransformImpl(tomo.backends.AstraCpuRayTransformImpl):
        def call_forward(...):
            my_custom_stuff()

    ray_trafo = odl.tomo.RayTransform(reco_space, geometry, impl=MyRayTransformImpl)

Any implementation must subclass RayTransformImplBase, essentially forcing the call_forward() and call_backward(). There are some optional methods in the base class, such as supports_geometry(geom) and supports_reco_space().

Behavior and API for RayTransform and RayBackProjection remains backward compatible, in particular you can still specify 'astra_cpu' etc. to impl. However, AstraCudaProjectorImpl and AstraCudaBackProjectorImpl are merged into AstraCudaRayTransformImpl. A lot of unnecessarily duplicated code has been removed there, and I didn't expect a lot of users to directly rely on these classes.

I'll make proper docs and fix the tests :) but I first wanted to send this in for some feedback on the concept.

@pep8speaks
Copy link

pep8speaks commented Jan 27, 2020

Checking updated PR...

No PEP8 issues.

Comment last updated at 2020-04-16 18:24:43 UTC

@adriaangraas adriaangraas force-pushed the raytrafo-strategy branch 2 times, most recently from d9fb3f1 to ca63f84 Compare January 28, 2020 20:37
Copy link
Member

@kohr-h kohr-h left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @adriaangraas. Thanks a lot for your contribution! I'm always in favor of making code more extensible and better structured. Your suggestion for a standardized ray transform interface is certainly interesting.

What I definitely like is the single class for forward and adjoint, since they share a lot of logic anyhow.

Maybe it's just me, but I've become a bit suspicious to class hierarchies recently, and this proposal solves the problem of unifying an interface by adding a base class layer to the implementations. That's a bit of a red flag to me and makes me wonder if there's no other way.
Regarding the extra methods like can_handle_size, I'm also not really sure if it's worth the trouble. Every additional method raises the threshold of implementing a subclass, even if the methods are optional. And if they're only used internally, why adding them? For instance, the check with the size warning could simply be in the __init__ of the implementation class. Who else would be interested?

I'm wondering what we want to achieve here? I guess the goal is to offer a point of extension (other than a PR) to users of the library who have their own ray transform implementation. It means that RayTransform should have some kind of lookup table for implementations that external users can amend at runtime. The basic tasks of an implementation are

  • initialization including sanity checks (default can be provided),
  • calls forward and backward,
  • (if necessary) managing of external objects (like ASTRA IDs)
  • (optional) caching

So in principle I'd expect that for an impl subclass, it would suffice to provide call_forward and call_backward. Would that be possible here?


Besides, I think we can easily scrap the RayTransformBase class in favor of an inline adjoint class. I know that was discussed at some point, and it was argued that RayBackProjection should be an operator on its own. But I've yet to see a case where it's insufficient or less intuitive to create a RayTransform and then take its adjoint.

odl/tomo/backends/impl.py Outdated Show resolved Hide resolved
odl/tomo/backends/impl.py Outdated Show resolved Hide resolved
odl/tomo/backends/impl.py Outdated Show resolved Hide resolved
odl/tomo/backends/skimage_radon.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/backends/astra_cpu.py Outdated Show resolved Hide resolved
@adriaangraas
Copy link
Contributor Author

adriaangraas commented Feb 16, 2020

Hi @kohr-h, thanks for reviewing!

I think you have fair points here. Although I like programming to interfaces (more to open the class for extension, not necessarily to standardize), I agree that OOP machinery easily ends up with a lot of boilerplate code. I tried removing the RayTransformImplBase, and I think that goes fairly well with just a few concessions: there is small bit of code duplication in the __init__ of each implementation, and I think it'd have been a good place to implement complex _calls. I added a few runtime checks in RayTransform verifying the existence of call_forward and/or call_backward. All in all I think it simplified the code quite a bit.

I also removed the RayTransformBase class, as you suggested. The RayTransform, and RayBackProjection are now directly derived from Operator. I completely agree with the former solution being suboptimal: subclassing, to have two variants, while simultaneously having the same variance in the forward and adjoint() is very confusing. Also, I don't think subclassing was the right pattern there to begin with. I'm not sure how you'd design the inner adjoint class, but I thought that a very natural way was to compose RayBackProjection using instead of extending a RayTransform. I gave this a try as well, let me know what you think. Another idea, by the way, is to write a mixin class to share logic between the two operators.

I've two direct questions:

  1. Now that I moved warnings into the implementations I cannot suppress the size warning, like here. I could either add a verbose flag, or try to catch the warnings, but neither seem very appealing. Do you have a good solution here?
  2. I used a complexify function to turn real-valued functions into complex-valued functions. Is there a nicer way, maybe using ODL utilities, to achieve this?

Copy link
Member

@kohr-h kohr-h left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the update @adriaangraas! I like the structure already a lot better than in the first iteration. Flattening the hierarchy definitely helped to clarify the logic. Some detail comments below (plus some nits about formatting and stuff 😉 ).

I'm not sure how you'd design the inner adjoint class, but I thought that a very natural way was to compose RayBackProjection using instead of extending a RayTransform. I gave this a try as well, let me know what you think.

A nice example of how this can work is here:

def gradient(self):
r"""Gradient operator of the functional.
Notes
-----
The derivative is computed using the quotient rule:
.. math::
[\nabla (f / g)](p) = (g(p) [\nabla f](p) -
f(p) [\nabla g](p)) / g(p)^2
"""
func = self
class FunctionalQuotientGradient(Operator):
"""Functional representing the gradient of ``f(.) / g(.)``."""
def _call(self, x):
"""Apply the functional to the given point."""
dividendx = func.dividend(x)
divisorx = func.divisor(x)
return ((1 / divisorx) * func.dividend.gradient(x) +
(- dividendx / divisorx**2) * func.divisor.gradient(x))
return FunctionalQuotientGradient(self.domain, self.domain,
linear=False)

The class is just defined inline, and the outer self is bound to a different name so that it's not shadowed. As a bonus, the inner class doesn't even need a dedicated __init__ since it can just use the one of the superclass Operator.

  1. Now that I moved warnings into the implementations I cannot suppress the size warning, like here. I could either add a verbose flag, or try to catch the warnings, but neither seem very appealing. Do you have a good solution here?

Did I mention that I don't really like the warnings at all?
One way I can think of would be to restrict the automatic selection to fast backends and force users to explicitly choose slow ones. In that case the warnings would be obsolete since slow backends are opt-in. @adler-j ?

  1. I used a complexify function to turn real-valued functions into complex-valued functions. Is there a nicer way, maybe using ODL utilities, to achieve this?

Well, my first thought was to translate your construct into a decorator for call_forward and call_backward, but that cannot work since the self argument is needed. What works instead is to make a wrapper function that handles both cases (I haven't tested the code):

def _add_default_complex_impl(fn):
    def wrapper(self, x, out, **kwargs):
        if self.reco_space.is_real and self.proj_space.is_real:
            fn(x, out, **kwargs)
            return out
        elif self.reco_space.is_complex and self.proj_space.is_complex:
            # NB: change calling code so that `out` cannot be `None`
            fn(x.real, out.real, **kwargs)
            fn(x.imag, out.imag, **kwargs)
            return out
        else:
            raise RuntimeError('mix of real and complex')

    return wrapper

@_add_default_complex_impl
def call_forward(self, x, out, **kwargs):
    # Code for real version goes here

@_add_default_complex_impl
def call_backward(self, x, out, **kwargs):
    # Code for real version goes here

odl/tomo/backends/astra_cpu.py Outdated Show resolved Hide resolved
odl/tomo/backends/astra_cpu.py Outdated Show resolved Hide resolved
odl/tomo/backends/astra_cpu.py Outdated Show resolved Hide resolved
odl/tomo/backends/astra_cpu.py Outdated Show resolved Hide resolved
"size. Consider using 'astra_cuda' if your machine has an "
"Nvidia GPU. This warning can be disabled by explicitly "
"setting `impl='astra_cpu'`.",
RuntimeWarning)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have to say that I never was a big fan of these warnings. They come from #1303, and you can read the discussion there for the reasoning.

odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
@adriaangraas
Copy link
Contributor Author

adriaangraas commented Feb 20, 2020

Thanks again @kohr-h, I don't mind the nitty comments 👍

Open remain mostly decisions from the unresolved conversations above. Is there anything else that would be good to add to this PR?

About the decorator: Although the syntax is really clean in the class I don't like that it implicitly passes self through the call (and then depends on self.vol_space and so on). I would have liked it better when it was more restricted to its input, or more widely applicable. I'll give it a thought.

In a next update I'll try to have a look at the docs and tests!

odl/tomo/backends/util.py Outdated Show resolved Hide resolved
@kohr-h
Copy link
Member

kohr-h commented Feb 20, 2020

Open remain mostly decisions from the unresolved conversations above. Is there anything else that would be good to add to this PR?

I'll go through it in detail again to get a better picture. But one thing that comes to my mind is the impl string. Maybe instead of using __class__.__name__, the ...Impl classes should have a class attribute impl_name or so, and register themselves with that name in a global dictionary in ray_trafo.py. That smells a bit like danger of circular imports, but I think if the code in ray_trafo.py knows nothing about the specific Impl's it shouldn't even have to import any of those modules.

Do you think that's reasonable?

(Or maybe that's already the current status.)

@adriaangraas
Copy link
Contributor Author

adriaangraas commented Feb 22, 2020

Hi @kohr-h. Hm, I don't know about auto-registering to the global dictionary, sounds like code that is hard to read and debug. I can see that this allows users to register their own class, which will then automatically be used by RayTransform. I think the expressive RayTransform(impl=MyImpl()) is less magical and achieves the same thing. There may not be a need for a priority-register, since a user will likely want to put their implementation as most important anyway. Furthermore RayTransform is already unaware of the default implementations, because it reads from _AVAILABLE_IMPLS. Maybe I am missing something here?

Update: I guess we could make _AVAILABLE_IMPLS and _IMPL_STR2TYPE public though, so that amending at runtime is possible.

@kohr-h
Copy link
Member

kohr-h commented Feb 25, 2020

Update: I guess we could make _AVAILABLE_IMPLS and _IMPL_STR2TYPE public though, so that amending at runtime is possible.

That's what I mean. Two things, then: maybe as a public dictionary, the name should be a bit more verbose. And we don't need the two above, _AVAILABLE_IMPLS could just be _IMPL_STR2TYPE.keys().

@adriaangraas
Copy link
Contributor Author

adriaangraas commented Feb 27, 2020

I've changed _IMPL_STR2TYPE to an OrderedDict named RAY_TRAFO_IMPLS. Maybe verboseness is good here when the variable goes public? I'm also fine with the other name. The last inserted item now automatically has the highest priority.
I didn't manage to remove the second global without introducing a lot of extra code, so I kept it in. RayTransform needs to have different lists in order to tell the user that an impl exists but is not installed. That list can remain private though.

@adriaangraas adriaangraas marked this pull request as ready for review February 27, 2020 21:51
Copy link
Member

@kohr-h kohr-h left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's been a while, but it's not forgotten. Just other stuff to do 🤷‍♂️

I'm quite happy with the structure and the implementation. The comments are for the final polish, and after that I think we're ready to merge the changes. Unless someone has objections. @adler-j? @aringh?

odl/tomo/backends/astra_cuda.py Outdated Show resolved Hide resolved
odl/tomo/backends/astra_cuda.py Outdated Show resolved Hide resolved
odl/tomo/backends/skimage_radon.py Outdated Show resolved Hide resolved
odl/tomo/backends/util.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Outdated Show resolved Hide resolved
@kohr-h
Copy link
Member

kohr-h commented Apr 13, 2020

@adriaangraas I'm currently rebasing on master. There has been a rather big internal API change (#1459) that impacts this PR, so I think it's easiest if I deploy the necessary changes here so you can focus on the actual content of the PR. Is that okay? So don't worry about an upcoming force-push 😉

@kohr-h kohr-h force-pushed the raytrafo-strategy branch from af3069d to f952b9e Compare April 13, 2020 12:03
kohr-h added 4 commits April 13, 2020 14:34
- Opening and closing parens on separate lines if line
  needs split
- No more backslash line continuation
- Sorted imports
- Removal of remaining references to `DiscreteLp`
@kohr-h
Copy link
Member

kohr-h commented Apr 13, 2020

Okay, I pushed a couple of commits: the rebase, fixes of rebase accidents and some formatting updates. Please take it from here @adriaangraas

Copy link
Member

@kohr-h kohr-h left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One little thing, and then we're good to go!

odl/tomo/backends/util.py Outdated Show resolved Hide resolved
odl/tomo/operators/ray_trafo.py Show resolved Hide resolved
@kohr-h
Copy link
Member

kohr-h commented Apr 16, 2020

One little thing

And I just did it myself :-)

@kohr-h
Copy link
Member

kohr-h commented Apr 16, 2020

Not waiting for CI. The last commit built fine, so I'll just merge it.

@kohr-h kohr-h merged commit 6fba0ca into odlgroup:master Apr 16, 2020
@kohr-h
Copy link
Member

kohr-h commented Apr 16, 2020

Thanks a lot for all the work @adriaangraas!

@adriaangraas
Copy link
Contributor Author

And thanks a ton for all the guidance @kohr-h, I've also learned a lot!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants