mims-harvard / shepherd Goto Github PK

def get_model(args, hparams, node_hparams, all_data, edge_attr_dict, n_nodes, load_from_checkpoint=False):
    print("setting up model", hparams['model_type'])
    # get patient model 
    if hparams['model_type'] == 'aligner':
        if load_from_checkpoint: 
            comb_patient_model = CombinedGPAligner.load_from_checkpoint(checkpoint_path=str(Path(project_config.PROJECT_DIR /  args.best_ckpt)), 
                                    edge_attr_dict=edge_attr_dict, all_data=all_data, n_nodes=n_nodes, node_ckpt = hparams["saved_checkpoint_path"], node_hparams=node_hparams)
        else:
            comb_patient_model = CombinedGPAligner(edge_attr_dict=edge_attr_dict, all_data=all_data, n_nodes=n_nodes, hparams=hparams, node_hparams=node_hparams)

class CombinedGPAligner(pl.LightningModule):

    def __init__(self, edge_attr_dict, all_data, n_nodes=None, node_ckpt = None, hparams=None, node_hparams=None,  spl_pca=[], spl_gate=[]):
        super().__init__()
        print('Initializing Model')

        self.save_hyperparameters('hparams', ignore=["spl_pca", "spl_gate"]) # spl_pca and spl_gate never get used

        print("Node checkpoint:", node_ckpt)

        print('Saved combined model hyperparameters: ', self.hparams)

        self.all_data = all_data

        self.all_train_nodes = {}
        self.train_patient_nodes = {}
        self.train_sparse_nodes = {}
        self.train_target_batch = {}
        self.train_corr_gene_nid = {}

        #print(f"Loading Node Embedder from {self.hparams.hparams['saved_checkpoint_path']}")
        print(f"Loading Node Embedder from {node_ckpt}")

        # NOTE: loads in saved hyperparameters
        self.node_model = NodeEmbeder.load_from_checkpoint(checkpoint_path=node_ckpt, #self.hparams.hparams['saved_checkpoint_path'], 
                                                           all_data=all_data, edge_attr_dict=edge_attr_dict, 
                                                           num_nodes=n_nodes)
                                                           #num_nodes=n_nodes, combined_training=self.hparams.hparams['combined_training']) ```

with open(str(project_config.PROJECT_DIR / 'patients' / spl_index_fname), 'wb') as handle:
        pickle.dump(spl_indexing, handle, protocol=pickle.HIGHEST_PROTOCOL)
np.save(str(project_config.PROJECT_DIR / 'patients' / spl_matrix_fname), patients_spl_matrix)

# sample nodes to form positive edges. we will try to predict these edges
row, col, e_id = self.adj_t_sample.coo()
# NOTE: only does self loops when no edges in the current partition of the dataset
target_batch = random_walk(row, col, source_batch, walk_length=1, coalesced=False)[:, 1]
batch = torch.cat([source_batch, target_batch], dim=0)


batch_size: int = len(batch)
        adjs = []
        n_id = batch
        for size in self.sizes:
            adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False)
.....