Som V Tambe

Julia Summer of Code Final Report (Blog)

Finally the time has come! It has been an adventurous 3 months, and a very different experience from what I had originally expected. I will share all the nitty gritties below. We implemented the project as a separate library known as DiffImages.jl. You can 🌟 it on GitHub!

PRs merged/accepted

  1. Add basic functions for colorspace transforms

  2. Add gradients for FilledExtrapolations

  3. Make gradient fixes

  4. Add rrules for Interpolations types and tests

  5. Add homography struct and make warp pipeline differentiable

  6. Add support for CoordinateTransformations transforms

  7. Add docs and demo

  8. Make ImageFiltering.imfilter pipeline differentiable

Long Term Goals

  1. Make imfilter completely differentiable and add examples of training kernels both as an implicit as well as explicit parameters in the documentation.

  2. Add documentation for how to create your own custom warps (non-linear, neural network, etc).

  3. Try to make composable transforms differentiable (ex. trfm = t1∘t2t_1 \circ t_2).

  4. Move all the adjoints to their respective upstream packages, like ImageTransformations, ImageFiltering, etc.

  5. Add more examples of things like training local homographies to the documentation.

  6. Try implementing the derivative βˆ‚L/βˆ‚X\partial L / \partial X.

Background

JuliaImages ecosystem hosts the major packages for image processing. It contains state-of-the-art implementations for image processing algorithms, and Julia itself provides a high level interface as a language which compiles to extremely machine code, which makes it as fast as C.

Meanwhile, Zygote extends automatic differentiation support to the Julia language. It is the main AD backend used in Julia's machine learning stack, Flux. Zygote provides differentiability support for any code that is Julia, which is why it has a much wider scope of usage apart from machine learning, like differential equations, density-functional theory, graphs, etc.

This project involved extending automatic differentiation to the JuliaImages ecosystem. We target various modules like warps, filters, colorspace transforms inside constituent packages of the ecosystem, such as ImageTransformations and ImageFiltering.

Motivation

JuliaImages consists of image processing libraries. Since there is a large amount of structural information that image processing modules carry, this information could be directly made use of without trying to train neural networks for the same thing.

Some very simple examples of these are training the parameters of a homography matrix, training local homographies for a large image, which could be used to stitch multiple images into one to form a wide angle image or a panorama, using manifolds as maps, and many more.

We aim to help enable these features using the general libraries in JuliaImages, and then finally push all the changes upstream so that users are able to differentiate through them directly without facing any hassles.

Progress this summer

This project turned out to be more of an open-ended problem than a closed-end problem which was what I had initially expected. I was asked to deviate from my proposed timeline as my mentors wanted me to make the warping pipeline differentiable. Ironically, I had put warps at the very end of my implementations, and because of the change, I had to review the literature and look at the existing libraries implementing warps.

My mentors were apt in their suggestions, as after we got the warping pipelines differentiable, various other things like filters came in line. If you think conceptually, a convolution or a correlation operation includes convolving over an image using a kernel (otherwise known as filters, as there can be other types too). This is why I took relatively much lesser time in writing the imfilter adjoints.

I shall explain my whole progress in chronological order:

We wanted to make a Flux-agnostic system which causes no problems in integrating any colorspace transform in a Chain(...) pipeline. Therefore I created two functions colorify and channelify which were fully differentiable (#3). These functions help us handle Images as batches using the original ColorTypes. We can now do something like this and get the gradients correctly -

julia> f = Chain(x->HSV.(x),channelify,flatten,Dense(768,16),Dense(16,10),x->Οƒ.(x))

julia> f(rand(RGB,16,16,1))
10Γ—1 Matrix{Float64}:
 1.659111580889885e-8
 1.0883817593664822e-185
 0.021796811490401094
 1.0
 8.61386828469479e-89
 2.1464761034716476e-60
 1.0
 3.2087902639230936e-48
 1.0
 1.0
julia> h = DiffImages.Homography{Float32}()
DiffImages.Homography{Float32} with:
3Γ—3 StaticArrays.SMatrix{3, 3, Float32, 9} with indices SOneTo(3)Γ—SOneTo(3):
 1.0  0.0  0.0
 0.0  1.0  0.0
 0.0  0.0  1.0

julia> h([1.0, 2.0, 3.0])
2-element StaticArrays.SVector{2, Float32} with indices SOneTo(2):
 0.33333334
 0.6666667

Using this, we could now warp homographies to images.

homog

(Training a rotation matrix involves training the parameters of a rotation matrix, which below are characterized simply by a coordinate space transformation by an angle of ΞΈ\theta. We train it against a target image which was rotated by an angle of Ο€/3\pi / 3, by starting at 0∘0^{\circ}. You can see how the original unrotated image slowly moves towards the angle of Ο€/3\pi / 3 after each iteration)

rot

This is still a work in progress, as I am currently working on this, and getting it to give correct gradients (#21).

Learnings

There were a ton of major learnings this summer. I firstly learned how to read large codebases (I am OKish at that for now), how to solve problems in the bottom up approach, and the temperament one has to deal with when having an open-ended problem at hand.

All these were very crucial, and there were times of frustration when I was debugging for days, not knowing what was going wrong. From those times, to finally being able to train and present something worthy of showcasing, I feel amazed at how much I have learnt over this JSoC period.

I was also exposed to a bunch of very smart people around me. In almost all (I think all) of the autodiff projects, the candidates were either PhD students or those who were going to pursue a PhD. Interacting with them at times made me feel at awe, and looking at the shear amount of knowledge they possessed motivated me to give my best.

Main part: my mentors

All this progress would not have been there had I not had such superb mentors. Dhairya Gandhi and Johnny Chen mentored me during my project duration. The insight and vision they had was really enlightening, and throughout all the stages, I have learnt and improved over a lot of aspects. I have learnt how to write good and meaningful code only because of the advice they have given me.

I also learnt to independently write code during this JSoC period. I have transitioned from writing ad hoc scripts to writing library level codes during the summer.

My mentors also taught me another very important thing - mentoring and teaching are two completely different things.

How tests make or break your project

One of the most important of writing library code is writing appropriate tests. My mentor Johnny always pressed on the importance of writing tests, and now that the summer is ending, I completely agree with what he wanted to convey in the very beginning.

Some final words

For me, JSoC was a completely unique experience I have ever had yet. I look forward to more opportunities which help me enhance my knowledge in the field of Computer Vision. I also learnt a lot about differentiable programming, and I'd love to explore more of it in the coming future in my free time.

I will also be experimenting with CUDA.jl and try adding CUDA support for Images. That should be another interesting task that I am eager to do.

Signing off,
Som ✌️