| 12 | class DiffusionEngine(core.Engine): |
| 13 | @torch.no_grad() |
| 14 | def generate(self, test_set, path): |
| 15 | if comm.get_rank() == 0: |
| 16 | logger.warning(f"Test on {test_set}") |
| 17 | path = os.path.expanduser(path) |
| 18 | if not os.path.exists(path): |
| 19 | os.makedirs(path) |
| 20 | logger.warning(path) |
| 21 | dataloader = data.DataLoader(test_set, self.batch_size, shuffle=False) |
| 22 | model = self.model |
| 23 | |
| 24 | model.eval() |
| 25 | id = 0 |
| 26 | data_dict = {} |
| 27 | for batch in dataloader: |
| 28 | if self.device.type == "cuda": |
| 29 | batch = utils.cuda(batch, device=self.device) |
| 30 | true_proteins = batch["graph"].clone() |
| 31 | pred_proteins = self.model.generate(batch)["graph"] |
| 32 | evaluation_metric = self.model.get_metric(pred_proteins, true_proteins, {}) |
| 33 | print(f"atom_rmsd_per_residue: {evaluation_metric['atom_rmsd_per_residue'].mean():<20}" |
| 34 | f"chi_0_mae_deg: {evaluation_metric['chi_0_ae_deg'].mean():<20}" |
| 35 | f"chi_1_mae_deg: {evaluation_metric['chi_1_ae_deg'].mean():<20}" |
| 36 | f"chi_2_mae_deg: {evaluation_metric['chi_2_ae_deg'].mean():<20}" |
| 37 | f"chi_3_mae_deg: {evaluation_metric['chi_3_ae_deg'].mean():<20}") |
| 38 | for p in pred_proteins.unpack(): |
| 39 | pdb_file = os.path.basename(test_set.pdb_files[id]) |
| 40 | protein = p.cpu() |
| 41 | protein.to_pdb(os.path.join(path, pdb_file)) |
| 42 | data_dict[pdb_file] = p.cpu() |
| 43 | id += 1 |