In Spectral Norm How to backprop through "u" and "v" as well

aSingh · April 21, 2019, 5:56am

It seems in the Spectral Norm –

pytorch/pytorch/blob/6d307db5b465b9679f8bcc6b2e075db5782b2906/torch/nn/utils/spectral_norm.py#L73


#    backproping through two forward passes, e.g., the common pattern in
#    GAN training: loss = D(real) - D(fake). Otherwise, engine will
#    complain that variables needed to do backward for the first forward
#    (i.e., the `u` and `v` vectors) are changed in the second forward.
weight = getattr(module, self.name + '_orig')
u = getattr(module, self.name + '_u')
v = getattr(module, self.name + '_v')
weight_mat = self.reshape_weight_to_matrix(weight)


if do_power_iteration:
    with torch.no_grad():
        for _ in range(self.n_power_iterations):
            # Spectral norm of weight equals to `u^T W v`, where `u` and `v`
            # are the first left and right singular vectors.
            # This power iteration produces approximations of `u` and `v`.
            v = normalize(torch.mv(weight_mat.t(), u), dim=0, eps=self.eps, out=v)
            u = normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps, out=u)
        if self.n_power_iterations > 0:
            # See above on why we need to clone
            u = u.clone()
            v = v.clone()

The “u” and “v” are treated as constants while doing backprop of loss with respect to weights (W).

Since the spectral norm
sigma(W) = u^T W v
then to estimate the derivatives of sigma(W) with respect to W, lets say I want to backprop through u and v too (as u and v are both functions of W). If I remove “with torch.no_grad()” will it compute the gradient of weights taking into account that “u” and "v"are also function of W ?