Method forward

graph_encoder.py:69–140

(self, g: dgl.DGLHeteroGraph)

        return pooled

    def forward(self, g: dgl.DGLHeteroGraph):
        device = next(self.parameters()).device
        g = g.to(device)
        try:
            from torch.amp import autocast
            from contextlib import nullcontext
            autocast_ctx = autocast(device_type='cuda', enabled=False) if device.type == 'cuda' else nullcontext()
        except Exception:
            from contextlib import ExitStack as nullcontext
            autocast_ctx = nullcontext()

        with autocast_ctx:
            self._ensure_type_embeddings(g, device)
            self._ensure_convs(g, device)

            try:
                total_edges = 0
                for et in g.canonical_etypes:
                    total_edges += g.num_edges(et)
            except Exception:
                total_edges = 0

            batch_size = self._infer_batch_size(g)
            if total_edges == 0:
                return torch.zeros(batch_size, self.out_dim, device=device, dtype=torch.float32)

            for ntype in g.ntypes:
                num_nodes = g.num_nodes(ntype)
                use_tid = 'tid' in g.nodes[ntype].data and num_nodes > 0
                if use_tid and ntype in ('diagnosis', 'descriptor'):
                    tid = g.nodes[ntype].data['tid'].to(device)
                    if ntype == 'diagnosis':
                        feat = self.node_norm(self.diag_emb(tid).float())
                    else:
                        feat = self.node_norm(self.desc_emb(tid).float())
                    g.nodes[ntype].data['h'] = torch.nan_to_num(feat, nan=0.0, posinf=0.0, neginf=0.0)
                else:
                    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
                    type_vec = self.type_embeddings[key]
                    h0 = type_vec.unsqueeze(0).expand(num_nodes, -1).contiguous()
                    g.nodes[ntype].data['h'] = h0.float()

            for layer_idx in range(self.n_layers):
                conv: HeteroGraphConv = self.convs[layer_idx]
                h = {ntype: g.nodes[ntype].data['h'] for ntype in g.ntypes}
                h = conv(g, h)
                for ntype in h:
                    h_nt = torch.relu(h[ntype]).float()
                    h_nt = self.dropout(h_nt)
                    g.nodes[ntype].data['h'] = torch.nan_to_num(h_nt, nan=0.0, posinf=0.0, neginf=0.0)

            pooled_list = []
            for ntype in g.ntypes:
                pooled = self._safe_mean_nodes(g, ntype)