Closed codingbutstillalive closed 1 year ago
Good question, I was also a bit puzzled by this at first sight of the code by I think here's what's going on:
linear_loss
and a cluster_loss
. both of these losses now have their own computation graph: the graph of linear_loss
depends on the learnable parameters of the linear layer and, in turn, the cluster_loss
depends only on the cluster parameters. loss
which already has its own, separate computation graph. the cool thing now is that these three separate computation graphs stay "in tact" after the +
operation and a grad_fn=<AddBackward0>
is added to the computation graph (more details here). e.g., if you print out the loss
variable you'll see tensor(3.3079, grad_fn=<AddBackward0>)
. self.manual_backward(loss)
here the autograd engine can track back which parts of the overall system the gradients came from and calculate the backward pass for the entire graph. the gradients of the linear and cluster parameters hence don't influence the updates of the weights of the segmentation head.Let me know if this helps. Would be interested in hearing your thoughts on this
Thank you very much. Both for the efforts of explaining it so well in detail, as well as for the additional insights. I totally appreciate that. This explanation resolves my confusion. It is a nice piece of engineering, although a few more code comments could have helped before releasing the code, I suppose ;)
@axkoenig Thank you for that great description of the technique! Couldn't have said it better myself, and appreciate your and @codingbutstillalive's patience :).
Based on the paper, the linear probe loss shall not be used for Stego's model training. But as far as I can see, both the linear probe loss as well as the cluster loss contribute to the overall loss used in backpropagation. Please explain.