Freeze Mapping Network and affine transformation layer

[zhangsdly]

As we know, freeze the mapping network and affine transformation layer during fine-tuning phase to better preserve semantic. The official repository only supports freezeD. if add freeze M and A, I think it is useful to explore the trained models without unnecessary headaches.

[PDillis]

Sounds like a good idea actually. Will add the option --freezeM and --freezeI (so we freeze the input on both SGAN2 and SGAN3).

[PDillis]

Added in commit: f43205a. E.g., when transferring from AFHQ:

Generator                     Parameters  Buffers  Output shape        Datatype
---                           ---         ---      ---                 ---     
mapping.fc0                   -           262656   [4, 512]            float32 
mapping.fc1                   -           262656   [4, 512]            float32 
mapping                       -           512      [4, 16, 512]        float32 
synthesis.input.affine        2052        -        [4, 4]              float32 
synthesis.input               262144      1545     [4, 512, 36, 36]    float32 
synthesis.L0_36_512.affine    262656      -        [4, 512]            float32 
synthesis.L0_36_512           2359808     25       [4, 512, 36, 36]    float32 
...

--freezeM seems to work! Please test it and let me know if any bug might've slipped by! Still missing --freezeI, but I should have it tomorrow.

[zhangsdly]

--freezeM works when freezing first layers of the Mapping Network. One Questions : when updating G_ema, may need to judge whether ‘mapping’ or 'affine' is in name of G_ema.buffers. if not, G_ema.mapping will be updated by G.

for b_ema, b in zip(G_ema.named_buffers(), G.named_buffers()):
    if ('mapping' in b_ema[0]) or ('affine' in b_ema[0]):
        continue
    else:
        b_ema[1].copy_(b[1])

[PDillis]

True, I have to study it more carefully (I merely did the same as with --freezeD, but there's only one Discriminator being tracked during training).

I've been having issues with --freezeM when training --cfg=stylegan2. Specifically, if I freeze all layers, then some strange errors happen that I haven't delved too much into, and if I freeze some, then I get an GPU OOM (which is strange, as training the full network with the same settings runs without a hitch). On the other hand, I am able to freeze one or both layers of --cfg=stylegan3-{r,t}, so there may be other checks I need to do, along with your point above.