I was going to ask if XDiff will be compatible with some Cuda kernels. It would be really nice to have a common set of primitive operations that can be used by XDiff, Knet, or ReverseDiff, because than one can use these libraries interchangeably, which would be interesting I think. It is something that Mike Innes have envisioned in his JuliaCon talk.
Regarding the segmented operations, TensorFlow support unsorted_segment_sum, which is written in GPU and takes the total number of segments as an argument, which might makes things easier, and that the rest is written in python.
I have tried to find the implementation in tensorflow repository on github, but I have failed. It is mentioned here
/tensorflow/python/ops/math_ops.py
but I have not find a definition and I have to confess that I am not a TensorFlow guru.
In my library, I have implemented each operation (layer) to have a allocated buffer and after the end of the operation, I have returned view into this buffer. This was the way, how I have dealt with the efficiency.
Broadcasting. XDiff already fuses all such sequential operations into a single expression (even if they are not sequential in source code) and GPUArrays already can deal with it.
Matrix multiplication. GPUArrays already use cuBLAS to handle it.
Convolutions. cuDNN should cover it, though AFAIK there’s no actively developed wrapper for it.
All other functions (like segmented_sum). I believe the most efficient way would be to implement such functions separately for CPU and GPU (CUDAnative should be helpful here).
I haven’t tested it yet, but XDiff’s generated code should work out of the box for (1) and (2) with GPUArrays, (3) should be easy to add given a good wrapper for cuDNN, and (4) depends on the primitives like segmented_sum, but should be doable in general.
As far as I understand, Knet uses AutoGrad.jl which is a port of Python’s AutoGrad, and both seem to support only CPU. I don’t know much about internals of these libraries, but I guess this is because of the wide range of operations that pure automatic differentiation can support - not all of these operations are easy (or even possible) to translate to GPU.
On the other hand, symbolic differentiation libraries like Theano, TensorFlow or XDiff produce code that is easy to translate to GPU, but lack some of the AD features (e.g. loops and conditions). So I’m not sure it’s possible to make these libraries completely interchangeable.
But what we definitely can do is to provide a set of common functions (e.g. relu, conv2d, etc.) with efficient CPU and GPU implementations that other packages can use.
But in your case bags isn’t a single number, right? Can you provide a definition of bags type?
# funcs.jl
loss(W, b, x, y) = sum(abs2.(y - (W * x .+ b)))
# main.jl
using XDiff
include("funcs.jl")
x, y = randn(11,1000), randn(1,1000)
W = randn(1,size(x,1))
b = rand()
df = xdiff(loss; W=W, b=b, x=x, y=y)
result = df(W, b, x, y)
# or
mem = Dict() # buffers
result = df(W, b, x, y; mem=mem)
Changes:
I moved functions to a separate file (needed only if you run the code from REPL). xdiff tries to extract the source code of the function using Sugar.jl. Sugar.jl in turn first tries to find source file for a function and, if it fails, parses its AST. If the function is defined in REPL, it doesn’t have a source file, and the code recovered from AST is somewhat hard for XDiff to understand.Thus loading functions from file is preferable (should be partially fixed in future releases).
“Inlined” predict into loss. For the same reason as above, I’ve switched off recursive parsing of nested functions for now.
Split θ into W and b because XDiff doesn’t support indexing in vector notation. The primary reason is that separate arrays are much easier to process on GPU than indexing/slicing of a single array.
Removed size - not sure it’s differentiable and it’s not really useful here assuming your are going to minimize the loss.