Hi, when I start to train a simple version StyleGAN2 in PyTorch 1.0.1.
Since I do not want to use retain_graph = True, A Runtime Error Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time. raised here.
The code of mine is below, really appreciate for your help, thank you!
...
# G & D definition.
G = G_stylegan2(fmap_base=opts.fmap_base,
                    resolution=opts.resolution,
                    mapping_layers=opts.mapping_layers,
                    opts=opts,
                    return_dlatents=opts.return_latents)
D = D_stylegan2(fmap_base=opts.fmap_base,
                    resolution=opts.resolution,
                    structure='resnet')
...
Loss_D_list = [0.0]
Loss_G_list = [0.0]
softplus = torch.nn.Softplus()
...
# Training.
for epoch in range(0, 100):
     bar = tqdm(loader)
     for i, (real_img,) in enumerate(bar):

            real_img = real_img.to(opts.device)
            latents = torch.randn([real_img.size(0), 512]).to(opts.device)

            # =======================================================================================================
            #   (1) Update D network: D_logistic_r1(default)
            # =======================================================================================================
            # Compute adversarial loss toward discriminator
            real_logit = D(real_img)
            fake_img = G(latents)
            fake_logit = D(fake_img.detach())

            d_loss = softplus(fake_logit).mean()
            d_loss = d_loss + softplus(-real_logit).mean()

            # original
            r1_penalty = D_logistic_r1(real_img.detach(), D)
            d_loss = d_loss + r1_penalty
            # lite
            # d_loss = d_loss.mean()

            loss_D_list.append(d_loss.item())

            # Update discriminator
            optim_D.zero_grad()
            d_loss.backward()
            optim_D.step()


            # =======================================================================================================
            #   (2) Update G network: G_logistic_ns_pathreg(default)
            # =======================================================================================================
            if i % 3 == 0:
                G.zero_grad()
                fake_scores_out = D(fake_img)
                g_loss = softplus(-fake_scores_out).mean()

                loss_G_list.append(g_loss.item())

                # Update generator
                # Note: g_loss.backward(retain_graph=True) is ok, but this version is not I want!
                g_loss.backward()
                optim_G.step()