u

Author: magic3007

main_gan.py

@@ -93,6 +93,8 @@ def main(config):
     if config.mode == 'train':
         trainer.fit(model, datamodule=dm,
                     ckpt_path=config.resume_ckpt_path)
+        trainer.test(model, datamodule=dm,
+                    ckpt_path=config.resume_ckpt_path)
     elif config.mode == 'test':
         trainer.test(model, datamodule=dm,
                      ckpt_path=config.resume_ckpt_path)

mol_utils.py

@@ -355,9 +355,5 @@ def save_mol_img(mols, f_name='tmp.png', is_test=False):
                 a_smi = Chem.MolToSmiles(a_mol)
                 mol_graph = read_smiles(a_smi)
 
-                break
-
-                # if not is_test:
-                #     break
         except:
             continue

parse_config.py

@@ -32,6 +32,7 @@ def get_GAN_config():
 
     # Training or Testing Config
     parser.add_argument('--lambda_wgan', type=float, help='weight of the wgan loss')
+    parser.add_argument('--resume_ckpt_path', type=str, help='resume training from the checkpoint')
 
     # Result Config
     parser.add_argument('--result_dir', type=str, help='Directory to save results')

plmodule_gan.py

@@ -130,7 +130,6 @@ def gp_norm(y, x):
             return ((dydx.norm(2, dim=1) - 1) ** 2).mean()
 
         # Random weight term for interpolation between real and fake samples
-
         edge_alpha = torch.rand(real_edges.size(0), 1, 1, 1).type_as(real_edges).requires_grad_(False)
         node_alpha = edge_alpha.reshape(-1, 1, 1).requires_grad_(False)
         # Get random interpolation between real and fake samples
@@ -330,10 +329,16 @@ def configure_optimizers(self):
         self.opt_d = torch.optim.Adam(self.D.parameters(), lr=self.hparams.lr_d)
         self.opt_v = torch.optim.Adam(self.V.parameters(), lr=self.hparams.lr_v)
         return self.opt_g, self.opt_d, self.opt_v
-
-    def on_val_epoch_end(self):
+    
+    def _shared_on_eval_epoch_end(self):
         edges_logits, nodes_logits = self.G(self.sampled_img_z.to(self.device))
         mols = self.get_gen_mols(nodes_logits, edges_logits, self.hparams.post_method)
         # Saving molecule images.
         mol_f_name = os.path.join(self.hparams.img_dir, 'mol-{}.png'.format(self.current_epoch))
-        save_mol_img(mols, mol_f_name, is_test=self.hparams.mode == 'test')
+        save_mol_img(mols, mol_f_name, is_test=True)
+        
+    def on_val_epoch_end(self):
+        self._shared_on_eval_epoch_end()
+        
+    def on_test_epoch_end(self):
+        self._shared_on_eval_epoch_end()

requirements.txt

@@ -0,0 +1,6 @@
+numpy==1.21.5
+pysmiles==1.0.1
+pytorch_lightning==1.6.3
+rdkit==2009.Q1-1
+scikit_learn==1.1.1
+torch==1.11.0