Identify trainable layers

[vionwinnie]

Hi I am trying to apply StyleGAN-NADA Clip method onto OASIS as suggested by the author, can you help me to identify which layer of OASIS is trainable?

[edgarschnfld]

Hi,

All layers defined in the generator's init function are trainable:

OASIS/models/generator.py

Lines 7 to 23 in 8fc1cf0

	class OASIS_Generator(nn.Module):
	def __init__(self, opt):
	super().__init__()
	self.opt = opt
	sp_norm = norms.get_spectral_norm(opt)
	ch = opt.channels_G
	self.channels = [16*ch, 16*ch, 16*ch, 8*ch, 4*ch, 2*ch, 1*ch]
	self.init_W, self.init_H = self.compute_latent_vector_size(opt)
	self.conv_img = nn.Conv2d(self.channels[-1], 3, 3, padding=1)
	self.up = nn.Upsample(scale_factor=2)
	self.body = nn.ModuleList([])
	for i in range(len(self.channels)-1):
	self.body.append(ResnetBlock_with_SPADE(self.channels[i], self.channels[i+1], opt))
	if not self.opt.no_3dnoise:
	self.fc = nn.Conv2d(self.opt.semantic_nc + self.opt.z_dim, 16 * ch, 3, padding=1)
	else:
	self.fc = nn.Conv2d(self.opt.semantic_nc, 16 * ch, 3, padding=1)

Looking at the StyleGAN-NADA page (https://stylegan-nada.github.io/), I see that you need a G_train and G_frozen.
You can freeze the generator with a function like this

def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag

G_frozen = model.netG # or model.netEMA
requires_grad(G_frozen, False)

You can also use this approach to freeze individual layers. This would allow you to train only early layers, late layers, or only the SPADE normalization layers, etc.