@YaoshengFu This PR will fix the spectral norm bug you see.

I really always appreciate you guys.

I have one question as to my understanding.

u and v can/should be updated in a in-palce manner while weight_orig not.
This is why weight_orig is registered as an attribute.

Is this correct?

@crcrpar Almost! weight_orig is still a parameter, so updates to it should be done on the parallelized module (the one you pass over to DataParallel) via things like optimizers explicitly started by user. Even if it is not updated in-place, it is fine, because the next time DataParallel copies it to the devices, the new weight_orig will be get from the parallelized module, and broadcast over.

u and v, however, are a bit different because they are updated when the module is activated, and not within users' control. So it is our job to automatically update it. And to ensure that such update works with DataParallel, it needs to be done in-place.

@t-vi and @crcrpar It would be great if one or both of you could take a look at the SN changes and let me know if they look good :)

Sorry for my late response.

Changes are good but I have some questions about test codes.
So could you tell me if you have the time?

Thank you for your comments @crcrpar :)

I am still having a problem with spectral norm complaining about in place changes when followed by batch normalization and distributed across 2 or more GPUs using data parallel. This architecture arises in Big GAN (https://arxiv.org/abs/1809.11096). Interestingly this is not an issue with group normalization, instance normalization, or no normalization. It is also not an issue on a single GPU.

So to get this a bit sorted, how many cases do we have? with/without DP * eval/training * weight.requires_grad=True/False * ???.
I'd feel much more comfortable if I had a list of everything that is expected to work.

@zmurez even with this patch?

@t-vi Yeah, I think those 8 cases are all we need to consider.

@zmurez even with this patch?

I think so. It complained regardless of normalization and number of GPU's prior to adding this patch. However, I just copied the relevant lines into my own spectral_norm.py file instead of pulling the entire branch... So it is possible I missed something... but it seems unlikely since all the other cases are fixed.

Note, I am still using the latest stable release. To get this patch to work I also grabbed a copy of the normalize function, and implemented my on chain_matmul (a single if statement in this case with 3 matrices).

Does this bug exist for you? If not I guess I will have to consider updating to the unstable version.

Thanks

@zmurez It is possible that other changes are needed. Do you have a small repro script I can try on my build?

@zmurez It is possible that other changes are needed. Do you have a small repro script I can try on my build?

import torch
import torch.nn as nn

#from spectral_norm import spectral_norm       # my local spectral_norm patch
from torch.nn.utils import spectral_norm       # torch implementation

dim=5
batchsize=10
net = nn.DataParallel(nn.Sequential(
            #nn.Conv2d(dim, dim,1),
            nn.BatchNorm2d(dim),
            #nn.LeakyReLU(0.2, inplace=True),
            spectral_norm(nn.Conv2d(dim, dim,1)),
        )).cuda()
noise = torch.randn(batchsize, dim, 1, 1).cuda()
net.zero_grad()
out = net(noise).sum()
out.backward()
print('no bug')

@zmurez Thanks! This is actually a bug elsewhere. I have submitted a fix at #13594. With these two patches together, I can run your script successfully.

@t-vi @crcrpar Thank you for your reviews! I appreciate them. :)

It seems that when spectral norm is applied to a conv module, u is randomly initialized. Then v is computed such that the invariant holds.

1. This initial solve is very very slow (over tens of minutes) for large networks.
   a) Is this really necessary or is a random initialization ok?
   b) Can this solve be replaced by a sequence of power iterations?

2. If the weights of the conv module are later initialized using nn.init, I think this invariant will be broken?

@zmurez Good point. I'll think about ways to fix it.

For the initialization I think it should be not terribly important whether to init u or v randomly, so that could likely be changed. For loading the state dict: this only happens when there you load an "old version" state dict. How about warning about the performance and (potentially) offering a switch to forward re-compute u instead of solving for v. (I think it might be save to "do the right, slow thing" by default rather than the other way round.)

I agree with t-vi. Random initialization of v is not harmful.

@t-vi @crcrpar @zmurez Thank you. I have removed the solving part at initialization in the new commit :)