Hello GBI Team,
thanks for creating and maintaining this neat Matlab library!
I was facing some issues regarding the adjoint calculation in applyJacobianT_
of StackMap
:
function x = applyJacobianT_(this, y, v)
% Reimplemented from :class:`Map`
x = cell(this.numMaps,1);
for n = 1:this.numMaps
x{n} = this.alpha(n) .* this.mapsCell{n}.applyJacobianT(y,v);
end
x = cat(length(this.sizeout),x{:});
end
- Size related run-time failure.
- Functional issue: the forward pass at its core is broadcast -> per-channel-op-application -> concat; so the backward pass should be slice -> per-channel-op-transpose-application -> sum-reduction.
function x = applyJacobianT_(this, y, v)
% Reimplemented from :class:`Map`
% The adjoint op of a StackMap (:= apply & concat) is slice &
% apply & summation.
% Therefore we need to slice the upstream gradient along its
% last dimension. To be working with arbitrary dimension, we
% could e.g. use approaches like:
% https://stackoverflow.com/questions/19955653/matlab-last-dimension-access-on-ndimensions-matrix
% For now, we take the simple path just flatten it to 2-D,
% followed by another reshape op back to basic.
% TODO(cl): quick'n'dirty, to be tested for robustness!
full_dim = size(y);
plain_dim = full_dim(1:end - 1);
cat_dim = full_dim(end);
y_flattened = reshape(y, [], cat_dim);
x = cell(this.numMaps,1);
for n = 1:this.numMaps
y_n = reshape(y_flattened(:, n), plain_dim);
x{n} = this.alpha(n) .* this.mapsCell{n}.applyJacobianT(y_n,v);
end
x = cat(length(this.sizeout),x{:});
x = sum(x, length(this.sizeout));
end
What's your thoughts on a preferred solution? Let me know! Of course, this needs testing and cleaning-up...
Somewhat related question. Calculating the adjoint of StackMap * UpstreamOp
, the adjoint of UpsteamOp
gets executed multiple times, once for each channel of StackMap
. Any better solution than wrapping StackMap
functionality to first execute and sum all channel contributions before passing the gradient to the upstream node?
Best,
CL