load and fix wts

Author: bhushan23

2_Latent-Shading-Gen/main_full.py

@@ -40,6 +40,9 @@ def main():
                         help='read first n rows (default: -1)')
     parser.add_argument('--details', type=str, default=None,
                         help='Explaination of the run')
+    parser.add_argument('--load_pretrained_model', type=str, default='../pretrained/net_epoch_r5_5.pth',
+                        help='Pretrained model path')
+
     if ON_SERVER:
         parser.add_argument('--syn_data', type=str, default='/nfs/bigdisk/bsonawane/sfsnet_data/',
                         help='Synthetic Dataset path')
@@ -71,6 +74,7 @@ def main():
     epochs     = args.epochs
     model_dir  = args.load_model
     read_first = args.read_first
+    pretrained_model_dict = args.load_pretrained_model
 
 
     if read_first == -1:
@@ -98,12 +102,17 @@ def main():
     else:
         print('Initializing weights')
         sfs_net_model.apply(weights_init)
+        sfs_net_pretrained_dict = torch.load(pretrained_model_dict)
+        sfs_net_state_dict = sfs_net_model.state_dict()
+        load_model_from_pretrained(sfs_net_pretrained_dict, sfs_net_state_dict)
+        sfs_net_model.load_state_dict(sfs_net_state_dict)
+    sfs_net_model.fix_weights()
 
     os.system('mkdir -p {}'.format(args.log_dir))
     with open(args.log_dir+'/details.txt', 'w') as f:
         f.write(args.details)
 
-    wandb.watch(sfs_net_model)
+    # wandb.watch(sfs_net_model)
 
     # 1. Train on both Synthetic and Real (Celeba) dataset
     train(sfs_net_model, syn_data, celeba_data=celeba_data, read_first=read_first,\