gan separate trainign models

Author: bhushan23

6_Shading_Resdiual_GAN_Separate_Train/data_loading.py

@@ -119,25 +119,7 @@ def get_sfsnet_dataset(syn_dir=None, read_from_csv=None, read_celeba_csv=None, r
     face    = []
     depth   = []
 
-    if read_from_csv is None:
-        for img in sorted(glob.glob(syn_dir + '*/*_albedo_*')):
-            albedo.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_face_*')):
-            face.append(img)    
-
-        for img in sorted(glob.glob(syn_dir + '*/*_normal_*')):
-            normal.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_depth_*')):
-            depth.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_mask_*')):
-            mask.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_light_*.txt')):
-            sh.append(img)
-    else:
+    if read_from_csv is not None:
         df = pd.read_csv(read_from_csv)
         if read_first is not None and len(df) > read_first:
             df = df.sample(read_first, random_state=100)
@@ -155,17 +137,17 @@ def get_sfsnet_dataset(syn_dir=None, read_from_csv=None, read_celeba_csv=None, r
         for _, v in name_to_list.items():
             v[:] = [syn_dir + el for el in v]
 
-        # Merge Synthesized Celeba dataset for Psedo-Supervised training
-        if read_celeba_csv is not None:
-            df = pd.read_csv(read_celeba_csv)
-            if read_first is not None and len(df) > read_first:
-                df = df.sample(read_first, random_state=100)
-            albedo += list(df['albedo'])
-            face   += list(df['face'])
-            normal += list(df['normal'])
-            depth  += list(df['depth'])
-            mask   += list(df['mask'])
-            sh     += list(df['light'])
+    # Merge Synthesized Celeba dataset for Psedo-Supervised training
+    if read_celeba_csv is not None:
+        df = pd.read_csv(read_celeba_csv)
+        if read_first is not None and len(df) > read_first:
+            df = df.sample(read_first, random_state=100)
+        albedo += list(df['albedo'])
+        face   += list(df['face'])
+        normal += list(df['normal'])
+        depth  += list(df['depth'])
+        mask   += list(df['mask'])
+        sh     += list(df['light'])
 
     assert(len(albedo) == len(face) == len(normal) == len(depth) == len(mask) == len(sh))
     dataset_size = len(albedo)

6_Shading_Resdiual_GAN_Separate_Train/main_full.py

@@ -90,7 +90,7 @@ def main():
     # return 
 
     # Init WandB for logging
-    wandb.init(project='SfSNet-CelebA-GANLoss-Shading-Residual-PreTrained')
+    wandb.init(project='SfSNet-CelebA-GAN-Separate-Training-Shading-Residual-PreTrained')
     wandb.log({'lr':lr, 'weight decay': wt_decay})
 
     # Initialize models
@@ -101,6 +101,9 @@ def main():
 
     if use_cuda:
         sfs_net_model = sfs_net_model.cuda()
+        albedo_gen_model = albedo_gen_model.cuda()
+        albedo_dis_model = albedo_dis_model.cuda()
+    
 
     if model_dir is not None:
         sfs_net_model.load_state_dict(torch.load(model_dir + 'sfs_net_model.pkl'))
@@ -168,6 +171,7 @@ def main():
 
     # fix the weights of albedo gen model
     albedo_gen_model.fix_weights()
+    sfs_net_model.fix_new_weights()
 
     train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data, celeba_data=celeba_data, read_first=read_first,\
            batch_size=batch_size, num_epochs=epochs, log_path=log_dir+'Mix_Training/', use_cuda=use_cuda, wandb=wandb, \

6_Shading_Resdiual_GAN_Separate_Train/models.py

@@ -228,6 +228,9 @@ def forward(self, x):
         out = self.conv3(out)
         return out
 
+    def fix_weights(self):
+        dfs_freeze(self)
+
 class LightEstimator(nn.Module):
     """ Estimate lighting from normal, albedo and conv features
     """
@@ -333,6 +336,9 @@ def fix_weights(self):
         dfs_freeze(self.light_estimator_model)
         # Note that we are not freezing Albedo gen model
 
+    def fix_new_weights(self):
+        dfs_freeze(self.albedo_residual_model_2)
+
 
 # Use following to fix weights of the model
 # Ref - https://discuss.pytorch.org/t/how-the-pytorch-freeze-network-in-some-layers-only-the-rest-of-the-training/7088/15

6_Shading_Resdiual_GAN_Separate_Train/train.py

@@ -298,7 +298,7 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
 
     # Load Synthetic dataset
     train_dataset, val_dataset = get_sfsnet_dataset(syn_dir=syn_data+'train/', read_from_csv=syn_train_csv, read_celeba_csv=celeba_train_csv, read_first=read_first, validation_split=2)
-    test_dataset, _ = get_sfsnet_dataset(syn_dir=syn_data+'test/', read_from_csv=syn_test_csv, read_celeba_csv=celeba_test_csv, read_first=100, validation_split=0)
+    test_dataset, _ = get_sfsnet_dataset(syn_dir=syn_data+'test/', read_from_csv=None, read_celeba_csv=celeba_test_csv, read_first=100, validation_split=0)
 
     syn_train_dl  = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
     syn_val_dl    = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)
@@ -324,16 +324,6 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
     os.system('mkdir -p {}'.format(out_syn_images_dir + 'val/'))
     os.system('mkdir -p {}'.format(out_syn_images_dir + 'test/'))
 
-    # Create Generator and Discriminator
-    
-    if use_cuda:
-        albedo_gen_model = albedo_gen_model.cuda()
-        albedo_dis_model = albedo_dis_model.cuda()
-    
-    # Init gen and disc
-    albedo_gen_model.apply(weights_init)
-    albedo_dis_model.apply(weights_init)
-
     # Collect model parameters
     model_parameters = sfs_net_model.parameters()
     optimizer = torch.optim.Adam(model_parameters, lr=lr) #, weight_decay=wt_decay)
@@ -345,10 +335,6 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
     g_optimizer = torch.optim.Adam(albedo_gen_model.parameters(), lr=lr)
     d_optimizer = torch.optim.Adam(albedo_dis_model.parameters(), lr=lr)
 
-    if use_cuda:
-        albedo_loss = albedo_loss.cuda()
-        recon_loss  = recon_loss.cuda()
-
     lamda_recon  = 1
     lamda_albedo = 10
 
@@ -377,9 +363,7 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
 
             # Apply Mask on input image
             # face = applyMask(face, mask)
-            predicted_normal, albedo_features, predicted_sh, shading_residual = sfs_net_model(face)
-            optimizer.zero_grad()
-            # GAN Training
+                        # GAN Training
             valid = torch.ones(albedo.shape[0], requires_grad = False)
             fake = torch.zeros(albedo.shape[0], requires_grad = False)
 
@@ -400,6 +384,7 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
                 real_sample = real_sample.cuda()
 
             # GAN loss
+            predicted_normal, albedo_features, predicted_sh, shading_residual = sfs_net_model(face)
             fake_albedo = albedo_gen_model(albedo_features)
             pred_fake   = albedo_dis_model(fake_albedo)
             # print(pred_fake.shape, valid.shape)
@@ -610,15 +595,15 @@ def train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data, celeba_da
             wandb_log_images(wandb, real_sh_face, mask, 'Train Real SH Predicted Face', epoch, 'Train Real SH Predicted Face', path=file_name + '_real_sh_face.png')
             wandb_log_images(wandb, syn_face, mask, 'Train Real SH GT Face', epoch, 'Train Real SH GT Face', path=file_name + '_syn_gt_face.png')
 
-            v_total, v_albedo, v_recon = predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
+            v_total, v_albedo, v_recon = predict_sfsnet(sfs_net_model, albedo_gen_model, syn_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
                                                                          out_folder=out_syn_images_dir+'/val/', wandb=wandb)
             # wandb.log({log_prefix + 'Val Total loss': v_total, log_prefix + 'Val Albedo loss': v_albedo, log_prefix + 'Val Recon loss': v_recon})
             print('Val set results: Total Loss: {}, Albedo Loss: {},  Recon Loss: {}'.format(v_total, v_albedo, v_recon))
 
             # Model saving
             torch.save(sfs_net_model.state_dict(), model_checkpoint_dir + 'sfs_net_model.pkl')
         if epoch % 5 == 0:
-            t_total, t_albedo, t_recon = predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
+            t_total, t_albedo, t_recon = predict_sfsnet(sfs_net_model, albedo_gen_model, syn_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
                                                                         out_folder=out_syn_images_dir + '/test/', wandb=wandb, suffix='Test')
 
             # wandb.log({log_prefix+'Test Total loss': t_total, log_prefix+'Test Albedo loss': t_albedo, log_prefix+'Test Recon loss': t_recon})
@@ -759,7 +744,7 @@ def train_with_shading_loss(sfs_net_model, syn_data, celeba_data=None, read_firs
             wandb_log_images(wandb, real_sh_face, mask, 'Train Real SH Predicted Face', epoch, 'Train Real SH Predicted Face', path=file_name + '_real_sh_face.png')
             wandb_log_images(wandb, syn_face, mask, 'Train Real SH GT Face', epoch, 'Train Real SH GT Face', path=file_name + '_syn_gt_face.png')
 
-            v_total, v_albedo, v_recon = predict_sfsnet(sfs_net_model, syn_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
+            v_total, v_albedo, v_recon = predict_sfsnet(sfs_net_model, albedo_gen_model, syn_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
                                                                          out_folder=out_syn_images_dir+'/val/', wandb=wandb)
             wandb.log({log_prefix + 'Val Total loss': v_total, log_prefix + 'Val Albedo loss': v_albedo, log_prefix + 'Val Recon loss': v_recon})
 
@@ -769,7 +754,7 @@ def train_with_shading_loss(sfs_net_model, syn_data, celeba_data=None, read_firs
             # Model saving
             torch.save(sfs_net_model.state_dict(), model_checkpoint_dir + 'sfs_net_model.pkl')
         if epoch % 5 == 0:
-            t_total, t_albedo, t_recon = predict_sfsnet(sfs_net_model, syn_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
+            t_total, t_albedo, t_recon = predict_sfsnet(sfs_net_model, albedo_gen_model, syn_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
                                                                         out_folder=out_syn_images_dir + '/test/', wandb=wandb, suffix='Test')
 
             wandb.log({log_prefix+'Test Total loss': t_total, log_prefix+'Test Albedo loss': t_albedo, log_prefix+'Test Recon loss': t_recon})