How to get a dlgradient result for a blackbox function

I'm working on a tomography reconstruction code using a neural representation, which I'd like to train. After some debugging I discovered that:

(1) you can't use the dlarrays as input to the radon() function;

(2) dlgradient will not work with untraced variables, therefore all variables must be dlarrays;

(3) Using extractdata() to get the input of the radon() function converted to double will fail because dlgradient will lose its trace.

So now I have a conundrum. I need to compute that gradient. dlgradient() won't do the job. What's the best alternative way to do it and have the same format as adamupdate() requires?

See below what my loss function looks like. The way it's written now it will fail at the Sinogram=radon()... line, because Intensity is a dlarray(). As I previously said, if you fix that by extracting the data, then dlgradient() fails.

Thoughts would be appreciated.

function [loss, gradients]=modelLoss(net, XY, Mask, SinogramRef, theta)
    Intensity = predict(net,XY);
    Intensity = reshape(Intensity, size(Mask));
    Intensity(Mask)=0;

    Sinogram=radon(Intensity, theta);
    Sinogram=dlarray(Sinogram);

    loss = sum((Sinogram(~Mask)-SinogramRef(~Mask)).^2);
    loss = loss / sum(~Mask(:));

    gradients = dlgradient(loss,net.Learnables);
end

Edit:

As pointed out below by Matt J (thanks!), the radon transform can be automatically converted to a matrix multiplication through the file exchange entry: https://www.mathworks.com/matlabcentral/fileexchange/44669-convert-linear-functions-to-matrix-form?s_tid=srchtitle

However, there's a few caveats that make this a very slow implementation (although it does work). For neural net training, it looks like it is not necessarily ideal, so I'll post my performance numbers as well. First, the fixed loss function for me now looks like this:

function [loss,gradients]=modelLoss(net, XY, Mask, SinogramRef, Amatrix,theta)
    Intensity = predict(net,XY);
    Intensity = reshape(Intensity, size(Mask));
    Intensity(Mask)=0;
    
    Sinogram=Amatrix*Intensity(:);
    Sinogram=reshape(Sinogram,[],length(theta));

    SinogramMask=isnan(SinogramRef);
    loss = sum((Sinogram(~SinogramMask)-SinogramRef(~SinogramMask)).^2);
    loss = loss / sum(~Mask(:));

    gradients = dlgradient(loss,net.Learnables);
end

Where Amatrix is pre-calculated using func2mat:

Amatrix=func2mat(@(x) radonFc(x,size(occlusionSlice),theta),occlusionSlice); %forward radon transform
Amatrix=full(Amatrix);

Note the second line: full(Amatrix). That is required; because the deep learning toolbox will not work with sparse matrices. I.e., you can't do dlarray(Amatrix) if Amatrix is sparse. This is a big part of the slowdown because for the Radon transform that A matrix is really large. Also note I used the Amatrix based on the function radonFc:

function R=radonFc(X,sz,theta)
    X=reshape(X,sz);
    R=radon(X,theta);
end

which uses the actual Radon transform function (instead of what Matt proposed). The imrotate method can create differences in the size of the transformed image; plus the actual projection is different than simply rotating the image (there are sub-pixels used in the Radon transform and a slightly different interpolation scheme).

Now a very simple training loop would look like:

accfun = dlaccelerate(@modelLoss);
for iteration=1:maxIterations
    [loss, grad]= dlfeval(accfun,net,inputData,occlusionSlice, targetData, Amatrix, theta);
    [net,averageGrad,averageSqGrad] = adamupdate(net,grad,averageGrad,averageSqGrad,i);    
end

However, even with dlaccelerate, this function is extremely slow. The dlfeval function takes ~5 seconds to evaluate on a high-end laptop GPU (RTX 4060) and Amatrix, for a 112x112 image, has a size of 57865x12544 of full memory and takes up ~3GB of memory (which is why my GPU is bogged down with memory transfers).

As always, with machine learning things are resource hungry. This one, however, is likely not scalable. The alternative is to spend more human time to build the projection functions manually. Between computer time and human time, I'm going to let the computer work on this one =)

[Chuguang Pan]

The radon fucntion is not supported by the deep learning toolbox. I think if you want to include the radon fucntion in the loss function, maybe you can reproduce it with other functions supported by deep learning toolbox.

Answer 1

[Matt J]

radon() can be expressed as a matrix multiplication, and matrix multiplications are of course supported by dlarrays. One option for obtaining the matrix representation is using func2mat in this FEX download,

https://www.mathworks.com/matlabcentral/fileexchange/44669-convert-linear-functions-to-matrix-form?s_tid=srchtitle

It may take some time to compute the matrix, but you only have to do it once.

[Matt J]

The FEX submission also has a function called localfunc2mat, which can be used to accelerate matrix computation. The matrix A doesn't correspond precisely to radon(), but in my opinion, it is better.

xsiz=[256,256]; %input image size


R=kron(    speye(xsiz(1)) , ones(1,xsiz(2)) );  %row summation matrix

tic;
    clear A
    for i=360:-1:1
    
      fcn=@(x) imrotate(x,1-i,'crop'); %desired rotation operation
    
      [a,ysiz]=localfunc2mat(fcn, xsiz, 2);
      
      A{i,1}=R*a;
    
    end
    
    ysiz=[ysiz(1),numel(A)];
    A=cell2mat(A); %The final forward projection matrix
toc

Elapsed time is 122.932742 seconds.

%% Comaprision test

X=min(phantom(xsiz(1)),0.05);

Y1=radon(X,0:359);
Y2=reshape( A*X(:) , ysiz);

tiledlayout(1,2);
fcn=@(z) imshow(z(any(z,2),:),[]);
nexttile, fcn(Y1); title Radon
nexttile, fcn(Y2); title 'Matrix Multiplication'

●

[Fernando]

This is it. Incredibly slow; but that's it; it works. Can't complain with "it works" =)

So I was able to fix the code with this. I'll edit the question to reflect the code fix. Thank you for pointing out this function exists ^.^

[Matt J]

No, it's the same speed as the original radon() function. You only need to compute the matrix once.

>> tic; Y2=reshape( A*X(:) , ysiz); toc

Elapsed time is 0.039052 seconds.

>> tic; Y1=radon(X,0:359); toc

Elapsed time is 0.045699 seconds.

In fact, it should be easier for dlgradient to work with this, because with an explicit matrix, the gradient is just grad*z = A.'*z

[Fernando]

Matt, the matrix multiplication is fast in standard matlab, but the overhead from the deep learning toolbox is quite significant. Maybe you can check the edit on my original question. Each evaluation takes about 5 seconds because dlarray() requires a full matrix. So your A matrix is huge, which I believe requires a lot of tracing under the hood for the dlarray implementation to ensure the automatic differentiation is possible.

I'd love to hear new ideas if you have any, though! =)

[Fernando]

To further establish this is the problem:

The operation:

Sinogram=Amatrix*Intensity(:);

Takes 2.295179 seconds if Intensity is a dlarray(), and 0.101472 seconds if Intensity is a normal Matlab array.

Further, the dlgradient() function requires another 2.337362 seconds of runtime. So about 5 seconds for a dlfeval().

[Matt J]

I can't see from your original code why you would be converting Amatrix to a dlarray. There is no reason in principle why dlfeval cannot trace sparse matrix multiplications, as the example below illustrates.

A=sprand(50000,5000,10/5000);
x0=dlarray(rand(5000,1));

tic;
[fval0,g0]=dlfeval(  @(x) lossFcn(x,A)  , x0);
toc

Elapsed time is 0.452848 seconds.

function [fval,g]=lossFcn(x,A)

   fval=sum(A*x)^2;
   
   g=dlgradient(fval,x);

end

[Matt J]

The imrotate method can create differences in the size of the transformed image;

Nope. We used the 'crop' flag to avoid that.

plus the actual projection is different than simply rotating the image (there are sub-pixels used in the Radon transform and a slightly different interpolation scheme).

That is true of Matlab's implementation of radon(), but the interpolation scheme is not a good one. It produces artifacts at certain angles, as demonstrated here.

But if you really wanted to subpixelize, you could modify the black box function given to localfunc2mat(), e.g.,

R=kron(    kron( speye(xsiz(1)) , [1;1]/2)' , ones(1,2*xsiz(2)) );  %row summation matrix

fcn=@(x) imrotate(repelem(x,2,2),1-i,'crop');

[Fernando]

You see, Matt, though I appreciate your minimum working examples; it just doesn't work in the integrated code. I get this when A is sparse:

Error using  * 
MTIMES (*) is not supported for one sparse argument and one single argument.

Error in deep.internal.recording.operations.MtimesOp/forward (line 25)
x = x*y;

Error in  *  (line 129)
ydata = xdata * ydata;

The problem seems to be that the x in A*x comes from predict(), so I think there's something there.

[Matt J]

As the error message says, matrix multiplication is not supported for multiplication of a sparse matrix with a single precision matrix. But that is nothing specific to the Deep Learning Toolbox and dlarrays. That is true in general:

 speye(2)*ones(2,'single')

Error using *
MTIMES (*) is not supported for one sparse argument and one single argument.

Just convert to double floats for the purposes of the loss and gradient calculation.

Sinogram=Amatrix*double( Intensity(:));

...

gradients = single( dlgradient(loss,net.Learnables) );

[Fernando]

Thanks, Matt. Looks like this did the trick for CPU training. Thanks for your help =)

Answer 2

[Joss Knight]

Hi Fernando. The hardcore solution to your problem (after using Matt's solution) is to implement the sparse matrix multiply function yourself using a DifferentiableFunction. This can be used in your dlarray code:

classdef SparseMatrixVectorMultiply < deep.DifferentiableFunction

    properties
        SparseMatrix
    end

    methods
        function fcn = SparseMatrixVectorMultiply(sparseMatrix)
            % Constructor
            [e-mail](1, ...
                SaveInputsForBackward=false, ...
                SaveOutputsForBackward=false, ...
                NumMemoryValues=0);
            fcn.SparseMatrix = sparseMatrix;
        end

        function [Y, memory] = forward(fcn, X)
            % Forward pass: Y = A * X
            Y = fcn.SparseMatrix * X;
            memory = [];
        end

        function dLdX = backward(fcn, dLdY, memory)
            % Backward pass: dLdX = A' * dLdY
            dLdX = fcn.SparseMatrix' * dLdY;
        end
    end
end

Call it inside your modelLoss function like this:

spmv = SparseMatrixVectorMultiply(AMatrix);
Sinogram = spmv(Intensity(:));

Where Amatrix is the original sparse version of the matrix.

[Matt J]

But I wonder what the advantage of this would be? As shown earlier, ordinary automatic differentation of dlarrays already accomplishes this.

Also, I think you would again need,

Y = fcn.SparseMatrix * double(X);

dLdX = fcn.SparseMatrix' * double(dLdY);

[Matt J]

Perhaps because it would allow restoring the output to the same type is the input?

        function [Y, memory] = forward(fcn, X)
            % Forward pass: Y = A * X
            
            c=underlyingType(X);
            Y = feval(  fcn.SparseMatrix * double(X));
            memory = [];
        end
        function dLdX = backward(fcn, dLdY, memory)
            % Backward pass: dLdX = A' * dLdY
            
                c= underlyingType(dLdY);
            dLdY = feval(  fcn.SparseMatrix' * double(dLdY));

        end

[Joss Knight]

Well, the advantage is that you are doing a sparse matrix-vector multiplication instead of a dense one. This saves memory and (according to the OP) is much faster, so should scale better.

You're right, I was being lazy with the types here. But I'd use cast like:

function [Y, memory] = forward(fcn, X)
    % Forward pass: Y = A * X
    Y = cast(fcn.SparseMatrix * double(X), 'like', X);
    memory = [];
end
function dLdX = backward(fcn, dLdY, memory)
    % Backward pass: dLdX = A' * dLdY
    dLdX = cast(fcn.SparseMatrix' * double(dLdY), 'like', dLdY);
end

This ensures the double-precision behaviour is hidden inside the operation. This will be necessary until MATLAB supports single precision sparse.

[Matt J]

Well, the advantage is that you are doing a sparse matrix-vector multiplication instead of a dense one.

Sure, but it was demonstrated in this comment that you do not need a DifferentiableFunction to handle sparse matrix multiplications. The ordinary automatic differentation engine will handle it.

This will be necessary until MATLAB supports single precision sparse.

Are you hinting that it is on it's way?! That would be very exciting news. People have been asking for it for decades.

[Joss Knight]

Sorry Matt, your link didn't work for me, which comment? All I see is Fernando saying that he has to make Amatrix full for it to work with dlarray and that makes it very slow. I mean, as a developer of the toolbox I know that radon and sparse operations are not supported by dlarray so I'm not sure how it could be working out of the box, but I could be missing something.

[Joss Knight]

Oh no, apologies, you're right, it seems they snuck that in there without me noticing: as long as Amatrix is not a dlarray, it will behave correctly.

>> dlfeval(@(A,x)dlgradient(sum(A*x,'all'),x), sprand(10, 10, 0.4), dlarray(randn(10,1)))

ans = 

  10×1 dlarray

    2.9293
    0.9959
    0.8366
    1.7619
    2.2032
    0.8939
    2.7909
    2.3375
    1.2663
    0.9108