Class GraphEncoder

graph_encoder.py:9–140

warnings.filterwarnings('ignore')

class GraphEncoder(nn.Module):


    def __init__(self, out_dim: int, args=None, hidden: int = 128, n_layers: int = 2, dropout: float = 0.1):
        super().__init__()
        self.out_dim = out_dim
        self.hidden = hidden
        self.n_layers = n_layers
        self.dropout = nn.Dropout(p=dropout)

        self.type_embeddings = nn.ParameterDict()

        self.diag_emb = nn.Embedding(len(DIAGNOSIS_VOCAB) + 1, hidden)
        self.desc_emb = nn.Embedding(len(DESCRIPTOR_VOCAB) + 1, hidden)
        self.node_norm = nn.LayerNorm(hidden)

        self.convs = nn.ModuleList([
            HeteroGraphConv({}, aggregate='sum') for _ in range(n_layers)
        ])

        self.proj = nn.Linear(hidden, out_dim)

    def _ensure_type_embeddings(self, g: dgl.DGLHeteroGraph, device):
        for ntype in g.ntypes:
            key = f"emb_{ntype}"
            if key not in self.type_embeddings:
                param = nn.Parameter(torch.zeros(self.hidden, device=device, dtype=torch.float32))
                nn.init.xavier_uniform_(param.unsqueeze(0))
                self.type_embeddings[key] = param

    def _ensure_convs(self, g: dgl.DGLHeteroGraph, device):
        for layer_idx in range(self.n_layers):
            conv: HeteroGraphConv = self.convs[layer_idx]
            if len(conv.mods) == 0:
                mods = {}
                for srctype, etype, dsttype in g.canonical_etypes:
                    mods[etype] = GraphConv(in_feats=self.hidden, out_feats=self.hidden, norm='both', allow_zero_in_degree=True)
                conv.mods = nn.ModuleDict(mods)
                self.convs[layer_idx] = conv.to(device)

    def _infer_batch_size(self, g: dgl.DGLHeteroGraph) -> int:
        if hasattr(g, 'batch_size'):
            return g.batch_size
        try:
            if len(g.ntypes) > 0 and hasattr(g, 'batch_num_nodes'):
                base_ntype = g.ntypes[0]
                return len(g.batch_num_nodes(base_ntype))
        except Exception:
            pass
        return 1

    def _safe_mean_nodes(self, g: dgl.DGLHeteroGraph, ntype: str):
        if ntype not in g.ntypes or g.num_nodes(ntype) == 0:
            return None
        pooled = dgl.mean_nodes(g, 'h', ntype=ntype)
        if pooled is None:
            return None
        pooled = torch.nan_to_num(pooled, nan=0.0, posinf=0.0, neginf=0.0).float()