Add support for CoordinateTransformations transforms #19
Conversation
- Added RotMatrix, SMatrix adjoints - Removed Homography adjoint
- Add tests for SMatrix, RotMatrix, LinearMap - Made changes to Homography tests to check for the types - Add breaking tests for RotMatrix vs. non-RotMatrix gradients
Codecov Report
@@ Coverage Diff @@
## main #19 +/- ##
==========================================
+ Coverage 82.35% 88.23% +5.88%
==========================================
Files 3 3
Lines 102 102
==========================================
+ Hits 84 90 +6
+ Misses 18 12 -6
Continue to review full report at Codecov.
|
There was a problem hiding this comment.
With the test passing, nothing too suspicious to me.
the gradients come in a nested form compared to the ones when I take using an AffineMap composed of purely SArrays.
I believe this is expected? If you try to get gradients for any neural network, you get something similar. @DhairyaLGandhi should know more about how to process this. I guess it's time to introduce Functors.jl to handle nested structures.
| _, ∇ϕ = rrule_via_ad(Zygote.ZygoteRuleConfig(), tform, SVector(p.I)) | ||
| ∇h, _ = ∇ϕ(∇τ) | ||
| ∇tform += Tangent{DiffImages.Homography}(H = ∇h.H) | ||
| ∇tform += Tangent{typeof(tform)}(;NamedTuple{flds}(∇h)...) |
There was a problem hiding this comment.
| ∇tform += Tangent{typeof(tform)}(;NamedTuple{flds}(∇h)...) | |
| ∇tform += Tangent{typeof(tform)}(NamedTuple{flds}(∇h)...) |
There was a problem hiding this comment.
Can't do this xD. I wanted to pass the field names too, so have to pass it as kwargs :p
If I splat the namedtuple directly, the values get passed but not the keys.
| return y, smat_const_pb | ||
| end | ||
|
|
||
| function ChainRulesCore.rrule(p::Type{RotMatrix{N, T, L}}, x) where {N, T, L} |
There was a problem hiding this comment.
Seems like this would be hard to scale to other kinds of transforms if we were to write rules for every transform individually. Are there a countable few of these? Do all of them need special handling or can we AD most? Do these compose to produce other transforms?
There was a problem hiding this comment.
Most transforms can be written directly using an AffineMap. Others, we can use some existing transform to make that transform. This means we (hopefully) won't have to write any more rules.
We can AD through most, and yes I think we can compose to produce other transforms too.
I was planning the composing thing in the next PR.
There was a problem hiding this comment.
It would be good to have a poc that demos what you're talking about.
There was a problem hiding this comment.
Yes I'll prepare it with the docs.
|
Use Optimisers.jl along with FluxML/Functors.jl#1 to optimise this. We can do a quick call and I can walk you through structural optimisation. Otherwise make sure you have the |
| tfm1 = recenter(RotMatrix(π/8), center(img)) # using RotMatrix | ||
| tfm2 = AffineMap(SMatrix{2, 2, Float64, 4}([cos(π/8) -sin(π/8); sin(π/8) cos(π/8)]), tfm1.translation) | ||
|
|
||
| zy1 = Zygote.gradient(f, tfm1) |
There was a problem hiding this comment.
Oh brilliant, are the warp! rules working as expected now?
|
Since I have made the required changes, can we get this PR merged? |
There was a problem hiding this comment.
Looks good to me. I'll leave it to @DhairyaLGandhi since there seems one or two unresolved comments.
4 participants
Through this PR, I have added support for the following -
LinearMap,AffineMapandRotMatrixtransforms are now differentiable.Homographytransform to make it on the lines ofCoordinateTransformations.LinearMap. Now it is differentiable without the rule I had written.Some things that are breaking -
AffineMapcombined with aRotMatrix, something on the lines ofthe gradients come in a nested form compared to the ones when I take using an
AffineMapcomposed of purelySArrays.They are something like this -
I should fix this in another PR.
This PR should solve #17 .