Fix shading training

Author: bhushan23

6_Shading_Resdiual_GAN_Separate_Train/train.py

@@ -68,8 +68,8 @@ def predict_celeba(sfs_net_model, dl, train_epoch_num = 0,
     # return average loss over dataset
     return tloss / len_dl
 
-def predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, dl, gan_real_dl, train_epoch_num = 0,
-                    use_cuda = False, out_folder = None, wandb = None, suffix = 'Val'):
+def predict_sfsnet_gan(sfs_net_model, albedo_gen_model, albedo_dis_model, dl, gan_real_dl, train_epoch_num = 0,
+                    use_cuda = False, out_folder = None, wandb = None, suffix = 'GAN Val'):
 
     # debugging flag to dump image
     fix_bix_dump = 0
@@ -193,6 +193,95 @@ def predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, dl, gan_re
     # return average loss over dataset
     return tloss / len_dl, aloss / len_dl, rloss / len_dl, ganloss/len_dl, disloss / len_dl
 
+def predict_sfsnet(sfs_net_model, albedo_gen_model, dl, train_epoch_num = 0,
+                    use_cuda = False, out_folder = None, wandb = None, suffix = 'Val'):
+ 
+    # debugging flag to dump image
+    fix_bix_dump = 0
+
+    albedo_loss = nn.SmoothL1Loss() #nn.L1Loss()
+    recon_loss  = nn.SmoothL1Loss() #nn.L1Loss() 
+
+    lamda_recon  = 1
+    lamda_albedo = 1
+
+    if use_cuda:
+        albedo_loss = albedo_loss.cuda()
+        recon_loss  = recon_loss.cuda()
+
+    tloss = 0 # Total loss
+    aloss = 0 # Albedo loss
+    rloss = 0 # Reconstruction loss
+
+    for bix, data in enumerate(dl):
+        albedo, normal, mask, sh, face = data
+        if use_cuda:
+            albedo = albedo.cuda()
+            normal = normal.cuda()
+            mask   = mask.cuda()
+            sh     = sh.cuda()
+            face   = face.cuda()
+
+        # Apply Mask on input image
+        # face = applyMask(face, mask)
+        # predicted_face == reconstruction
+        # predicted_normal, predicted_albedo, predicted_sh, predicted_shading, shading_residual, updated_shading, predicted_face = sfs_net_model(face)
+
+        # Apply Mask on input image
+        # face = applyMask(face, mask)
+        predicted_normal, albedo_features, predicted_sh, shading_residual = sfs_net_model(face)
+
+        fake_albedo = albedo_gen_model(albedo_features)
+
+        out_shading = get_shading(predicted_normal, predicted_sh)
+        updated_shading = out_shading + shading_residual
+        out_recon = reconstruct_image(updated_shading, fake_albedo)
+
+        # albedo recon loss
+        current_albedo_loss = albedo_loss(fake_albedo, albedo)
+        current_recon_loss  = recon_loss(out_recon, face)
+
+        total_loss = lamda_albedo * current_albedo_loss + lamda_recon * current_recon_loss 
+
+        # Logging for display and debugging purposes
+        tloss += total_loss.item()
+        # nloss += current_normal_loss.item()
+        aloss += current_albedo_loss.item()
+        # shloss += current_sh_loss.item()
+        rloss += current_recon_loss.item()
+
+        if bix == fix_bix_dump:
+            # save predictions in log folder
+            file_name = out_folder + suffix + '_' + str(train_epoch_num) + '_' + str(fix_bix_dump)
+            # log images
+            # save_p_normal = get_normal_in_range(predicted_normal)
+            save_gt_normal = get_normal_in_range(normal)
+            save_p_normal = predicted_normal
+
+            wandb_log_images(wandb, save_p_normal, mask, suffix+' Predicted Normal', train_epoch_num, suffix+' Predicted Normal', path=file_name + '_predicted_normal.png')
+            wandb_log_images(wandb, fake_albedo, mask, suffix +' Predicted Albedo', train_epoch_num, suffix+' Predicted Albedo', path=file_name + '_predicted_albedo.png')
+            wandb_log_images(wandb, out_shading, mask, suffix+' Predicted Shading', train_epoch_num, suffix+' Predicted Shading', path=file_name + '_predicted_shading.png', denormalize=False)
+            wandb_log_images(wandb, shading_residual, mask, suffix+' Predicted Shading Residual', train_epoch_num, suffix+' Predicted Shading Residual', path=file_name + '_predicted_residual_shading.png', denormalize=False)
+            wandb_log_images(wandb, updated_shading, mask, suffix+' Predicted Updated Shading', train_epoch_num, suffix+' Predicted Updated Shading', path=file_name + '_predicted_updated_shading.png', denormalize=False)
+            wandb_log_images(wandb, out_recon, mask, suffix+' Predicted face', train_epoch_num, suffix+' Predicted face', path=file_name + '_predicted_face.png', denormalize=False)
+            wandb_log_images(wandb, face, mask, suffix+' Ground Truth', train_epoch_num, suffix+' Ground Truth', path=file_name + '_gt_face.png')
+            wandb_log_images(wandb, save_gt_normal, mask, suffix+' Ground Truth Normal', train_epoch_num, suffix+' Ground Normal', path=file_name + '_gt_normal.png')
+            wandb_log_images(wandb, albedo, mask, suffix+' Ground Truth Albedo', train_epoch_num, suffix+' Ground Albedo', path=file_name + '_gt_albedo.png')
+            # Get face with real SH
+            real_sh_face = sfs_net_model.get_face(sh, predicted_normal, fake_albedo)
+            wandb_log_images(wandb, real_sh_face, mask, 'Val Real SH Predicted Face', train_epoch_num, 'Val Real SH Predicted Face', path=file_name + '_real_sh_face.png')
+            syn_face     = sfs_net_model.get_face(sh, normal, albedo)
+            wandb_log_images(wandb, syn_face, mask, 'Val Real SH GT Face', train_epoch_num, 'Val Real SH GT Face', path=file_name + '_syn_gt_face.png')
+
+            # TODO:
+            # Dump SH as CSV or TXT file
+    
+    len_dl = len(dl)
+    wandb.log({suffix+' Total loss': tloss/len_dl, suffix+'Albedo loss': aloss/len_dl, suffix + 'Recon loss': rloss/len_dl}, step=train_epoch_num)
+            
+    # return average loss over dataset
+    return tloss / len_dl, aloss / len_dl, rloss / len_dl
+
 def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data, celeba_data=None, read_first=None,
           batch_size = 10, num_epochs = 10, log_path = './results/metadata/', use_cuda=False, wandb=None,
           lr = 0.01, wt_decay=0.005):
@@ -342,7 +431,6 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
 
             loss_d.backward()
             d_optimizer.step()
-            optimizer.step()
 
             # Logging for display and debugging purposes
             tloss += total_loss.item()
@@ -384,7 +472,7 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
             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, v_gloss, v_dloss = predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_val_dl, gan_real_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
+            v_total, v_albedo, v_recon, v_gloss, v_dloss = predict_sfsnet_gan(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_val_dl, gan_real_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
                                                                          out_folder=out_syn_images_dir+'/val/', wandb=wandb, suffix='GAN Val')
             # wandb.log({log_prefix + 'Val Total loss': v_total, log_prefix + 'Val Albedo loss': v_albedo, log_prefix + 'Val Recon loss': v_recon})
 
@@ -394,11 +482,9 @@ def gan_based_train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data,
             # 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, t_gloss, t_dloss = predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_test_dl, gan_real_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
+            t_total, t_albedo, t_recon, t_gloss, t_dloss = predict_sfsnet_gan(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_test_dl, gan_real_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
                                                                         out_folder=out_syn_images_dir + '/test/', wandb=wandb, suffix='GAN Test')
 
-            # wandb.log({log_prefix+'Test Total loss': t_total, log_prefix+'Test Albedo loss': t_albedo, log_prefix+'Test Recon loss': t_recon})
-
             print('Test-set results: Total Loss: {}, Albedo Loss: {}, Gan Loss: {}, Dis Loss: {} \n'.format(t_total, t_albedo, t_gloss, t_dloss))
 
 def train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data, celeba_data=None, read_first=None,
@@ -423,15 +509,6 @@ def train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data, celeba_da
     syn_val_dl    = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)
     syn_test_dl   = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
 
-    gan_real_train_dataset, gan_real_val_dataset = get_sfsnet_dataset(syn_dir=syn_data+'train/', read_from_csv=syn_train_csv, read_celeba_csv=None, read_first=read_first, validation_split=2)
-    gan_real_test_dataset, _ = get_sfsnet_dataset(syn_dir=syn_data+'test/', read_from_csv=syn_test_csv, read_celeba_csv=None, read_first=read_first, validation_split=0)
-    
-    gan_real_train_dl  = DataLoader(gan_real_train_dataset, batch_size=batch_size, shuffle=True)
-    train_real_gan_iter = iter(gan_real_train_dl)
-
-    gan_real_val_dl     = DataLoader(gan_real_val_dataset, batch_size=batch_size, shuffle=True)
-    gan_real_test_dl  = DataLoader(gan_real_test_dataset, batch_size=batch_size, shuffle=True)
-
     print('Synthetic dataset: Train data: ', len(syn_train_dl), ' Val data: ', len(syn_val_dl), ' Test data: ', len(syn_test_dl))
 
     model_checkpoint_dir = log_path + 'checkpoints/'
@@ -448,32 +525,24 @@ def train(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_data, celeba_da
     optimizer = torch.optim.Adam(model_parameters, lr=lr) #, weight_decay=wt_decay)
     albedo_loss = nn.SmoothL1Loss() #nn.L1Loss()
     recon_loss  = nn.SmoothL1Loss() #nn.L1Loss() 
-    gan_loss = torch.nn.MSELoss()
-
-    # Collect and initialize gen-dis optimizers
-    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
+    lamda_albedo = 1
 
     if use_cuda:
         albedo_loss = albedo_loss.cuda()
         recon_loss  = recon_loss.cuda()
-        gan_loss    = gan_loss.cuda()
 
     syn_train_len    = len(syn_train_dl)
 
     for epoch in range(1, num_epochs+1):
         tloss = 0 # Total loss
         aloss = 0 # Albedo loss
         rloss = 0 # Reconstruction loss
-        ganloss = 0 # Gan Loss
-        disloss = 0 # Dis Loss
 
         for bix, data in enumerate(syn_train_dl):
             albedo, normal, mask, sh, face = data
@@ -541,15 +610,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, v_gloss, v_dloss = predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_val_dl, gan_real_val_dl, train_epoch_num=epoch, use_cuda=use_cuda,
+            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,
                                                                          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, t_gloss, t_dloss = predict_sfsnet(sfs_net_model, albedo_gen_model, albedo_dis_model, syn_test_dl, gan_real_test_dl, train_epoch_num=epoch, use_cuda=use_cuda, 
+            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, 
                                                                         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})