Update vtoonify/train_vtoonify_d.py

vtoonify/train_vtoonify_d.py

@@ -391,6 +391,7 @@ def train(args, generator, discriminator, g_optim, d_optim, g_ema, percept, pars
                 
                 
 
+
 if __name__ == "__main__":
     
     device = "cuda"
@@ -430,4 +431,86 @@ if __name__ == "__main__":
     if not args.pretrain:
         generator.encoder.load_state_dict(torch.load(args.encoder_path, map_location=lambda storage, loc: storage)["g_ema"])
         # we initialize the fusion modules to map f_G \otimes f_E to f_G.
-        fo
+        for k in generator.fusion_out:
+            k.conv.weight.data *= 0.01
+            k.conv.weight[:,0:k.conv.weight.shape[0],1,1].data += torch.eye(k.conv.weight.shape[0]).cuda()  
+        for k in generator.fusion_skip:
+            k.weight.data *= 0.01
+            k.weight[:,0:k.weight.shape[0],1,1].data += torch.eye(k.weight.shape[0]).cuda()
+
+    accumulate(g_ema.encoder, generator.encoder, 0)
+    accumulate(g_ema.fusion_out, generator.fusion_out, 0)
+    accumulate(g_ema.fusion_skip, generator.fusion_skip, 0) 
+
+    g_parameters = list(generator.encoder.parameters()) 
+    if not args.pretrain:
+        g_parameters = g_parameters + list(generator.fusion_out.parameters()) + list(generator.fusion_skip.parameters())
+
+    g_optim = optim.Adam(
+        g_parameters,
+        lr=args.lr,
+        betas=(0.9, 0.99),
+    )
+
+    if args.distributed:
+        generator = nn.parallel.DistributedDataParallel(
+            generator,
+            device_ids=[args.local_rank],
+            output_device=args.local_rank,
+            broadcast_buffers=False,
+            find_unused_parameters=True,
+        )
+
+    parsingpredictor = BiSeNet(n_classes=19)
+    parsingpredictor.load_state_dict(torch.load(args.faceparsing_path, map_location=lambda storage, loc: storage))
+    parsingpredictor.to(device).eval()
+    requires_grad(parsingpredictor, False)
+
+    # we apply gaussian blur to the images to avoid flickers caused during downsampling
+    down = Downsample(kernel=[1, 3, 3, 1], factor=2).to(device)
+    requires_grad(down, False)
+
+    directions = torch.tensor(np.load(args.direction_path)).to(device) 
+
+    # load style codes of DualStyleGAN
+    exstyles = np.load(args.exstyle_path, allow_pickle='TRUE').item()
+    if args.local_rank == 0 and not os.path.exists('checkpoint/%s/exstyle_code.npy'%(args.name)):
+        np.save('checkpoint/%s/exstyle_code.npy'%(args.name), exstyles, allow_pickle=True)    
+    styles = []
+    with torch.no_grad(): 
+        for stylename in exstyles.keys():
+            exstyle = torch.tensor(exstyles[stylename]).to(device)
+            exstyle = g_ema.zplus2wplus(exstyle)
+            styles += [exstyle]
+    styles = torch.cat(styles, dim=0)
+
+    if not args.pretrain:
+        discriminator = ConditionalDiscriminator(256, use_condition=True, style_num = styles.size(0)).to(device)
+
+        d_optim = optim.Adam(
+            discriminator.parameters(),
+            lr=args.lr,
+            betas=(0.9, 0.99),
+        )    
+
+        if args.distributed:
+            discriminator = nn.parallel.DistributedDataParallel(
+                discriminator,
+                device_ids=[args.local_rank],
+                output_device=args.local_rank,
+                broadcast_buffers=False,
+                find_unused_parameters=True,
+            )
+
+        percept = lpips.PerceptualLoss(model="net-lin", net="vgg", use_gpu=device.startswith("cuda"), gpu_ids=[args.local_rank])
+        requires_grad(percept.model.net, False)
+
+        pspencoder = load_psp_standalone(args.style_encoder_path, device)  
+
+    if args.local_rank == 0:
+        print('Load models and data successfully loaded!')
+
+    if args.pretrain:
+        pretrain(args, generator, g_optim, g_ema, parsingpredictor, down, directions, styles, device)
+    else:
+        train(args, generator, discriminator, g_optim, d_optim, g_ema, percept, parsingpredictor, down, pspencoder, directions, styles, device)