from args import get_args import torch.nn as nn import torch class BondWeightLayer(nn.Module): def __init__(self, bond_types=4, hidden_dim=64): super().__init__() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.bond_embedding = nn.Embedding(bond_types, hidden_dim) # Learnable bond representation self.edge_mlp = nn.Sequential( nn.Linear(hidden_dim, 1), ) self.edge_mlp = self.edge_mlp.to(device) # Move edge MLP to correct device nn.init.xavier_uniform_(self.bond_embedding.weight) # Xavier for embeddings for layer in self.edge_mlp: if isinstance(layer, nn.Linear): nn.init.xavier_uniform_(layer.weight) nn.init.zeros_(layer.bias) def forward(self, edge_types): bond_feats = self.bond_embedding(edge_types) # Convert bond type to learnable vector edge_weight = self.edge_mlp(bond_feats).squeeze() # Compute edge weight edge_weight = edge_weight.squeeze() # Compute edge weight return edge_weight def init_weights(m): if isinstance(m, nn.Linear): nn.init.kaiming_uniform_(m.weight, nonlinearity='leaky_relu') if m.bias is not None: nn.init.zeros_(m.bias) import torch.nn.functional as F from torch_geometric.nn import GINEConv import torch.nn as nn from vq import VectorQuantize # Ensure correct import class BondWeightLayer(nn.Module): def __init__(self, bond_types=4, hidden_dim=64): import torch super().__init__() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.bond_embedding = nn.Embedding(bond_types, hidden_dim) # Learnable bond representation self.edge_mlp = nn.Sequential( nn.Linear(hidden_dim, 1), nn.Sigmoid() # Output weight in range (0,1) ) self.edge_mlp = self.edge_mlp.to(device) # Move edge MLP to correct device def forward(self, edge_types): bond_feats = self.bond_embedding(edge_types) # Convert bond type to learnable vector edge_weight = self.edge_mlp(bond_feats).squeeze() # Compute edge weight return edge_weight import torch def _zeros_like(x, device=None, dtype=None): return torch.zeros((), device=device or getattr(x, "device", None), dtype=dtype or getattr(x, "dtype", None)) import torch def _normalize_quantize_output(qo, logger, device=None, dtype=None): """ Normalize quantizer outputs to 7-tuple: (loss, embed, commit_loss, cb_loss, sil_loss, repel_loss, cb_repel_loss) Assumption (ALWAYS): qo == (total_loss, (commit_loss, codebook_loss, repel_loss, cb_repel_loss)) Rules: - Never detach. - Never wrap existing tensors into new tensors. - Only convert python/np scalars/None to tensors. """ def _infer_device_dtype(*xs): nonlocal device, dtype if device is None or dtype is None: for x in xs: if torch.is_tensor(x): if device is None: device = x.device if dtype is None: dtype = x.dtype return if device is None: device = "cpu" if dtype is None: dtype = torch.float32 def _scalar_to_tensor(x): if torch.is_tensor(x): return x if x is None: _infer_device_dtype() return torch.zeros((), device=device, dtype=dtype) _infer_device_dtype() return torch.as_tensor(x, device=device, dtype=dtype) # --------------------------- # only supported shape: # (loss, (commit, cb, rep, cb_rep)) # --------------------------- if not (isinstance(qo, (tuple, list)) and len(qo) == 2 and isinstance(qo[1], (tuple, list))): raise TypeError(f"Expected (loss, (commit, cb, rep, cb_rep)) but got {type(qo)}: {qo}") loss, inner = qo if len(inner) != 5: raise TypeError(f"Expected inner tuple length 5: (commit, cb, rep, cb_rep, ent), got {len(inner)}: {inner}") commit, cb, rep, cb_rep, ent = inner[0], inner[1], inner[2], inner[3], inner[4] # total_loss, (commit_loss, codebook_loss, repel_loss, cb_repel_loss, ent_loss) logger.info(f"end of model forward, loss {loss},　commit {commit}, ent {ent}") _infer_device_dtype(loss, commit, cb, rep, cb_rep, ent) embed = None # qoにはembedが無い前提 sil = None # qoにはsilが無い前提（0で埋める） return ( _scalar_to_tensor(loss), embed, _scalar_to_tensor(commit), _scalar_to_tensor(cb), _scalar_to_tensor(sil), # silhouette_loss = 0 _scalar_to_tensor(rep), _scalar_to_tensor(cb_rep), _scalar_to_tensor(ent), ) import torch.nn as nn import torch import torch.nn as nn from utils import CORE_ELEMENTS class AtomEmbedding(nn.Module): def __init__(self): super(AtomEmbedding, self).__init__() # ---- 埋め込み定義（離散 0-30 部分）---- self.degree_embed = nn.Embedding(num_embeddings=7, embedding_dim=4) # 0..6 self.ring_embed = nn.Embedding(num_embeddings=2, embedding_dim=4) # ring flag+1 など self.charge_embed = nn.Embedding(num_embeddings=8, embedding_dim=4) # 0..7 にクリップ self.aromatic_embed = nn.Embedding(num_embeddings=2, embedding_dim=4) # 0/1 self.hybrid_embed = nn.Embedding(num_embeddings=6, embedding_dim=4) # 0..5 self.hydrogen_embed = nn.Embedding(num_embeddings=5, embedding_dim=4) # 0..4 ELEMENTS = [5, 6, 7, 8, 14, 15, 16] self.register_buffer( "element_lut", self._build_element_lut(ELEMENTS) ) self.element_embed = nn.Embedding(num_embeddings=len(ELEMENTS), embedding_dim=4) # 0/1 フラグ系は全部 2 クラス想定 def flag_emb(): return nn.Embedding(num_embeddings=2, embedding_dim=2) # 官能基フラグ 18 個 (7-24) self.func_embed_0 = flag_emb() self.func_embed_1 = flag_emb() self.func_embed_2 = flag_emb() self.func_embed_3 = flag_emb() self.func_embed_4 = flag_emb() self.func_embed_5 = flag_emb() self.func_embed_6 = flag_emb() self.func_embed_7 = flag_emb() self.func_embed_8 = flag_emb() self.func_embed_9 = flag_emb() self.func_embed_10 = flag_emb() self.func_embed_11 = flag_emb() self.func_embed_12 = flag_emb() self.func_embed_13 = flag_emb() self.func_embed_14 = flag_emb() self.func_embed_15 = flag_emb() self.func_embed_16 = flag_emb() self.func_embed_17 = flag_emb() # H-bond Donor / Acceptor (25, 26) self.h_don_embed = flag_emb() self.h_acc_embed = flag_emb() # ring size / aromatic neighbors / fused-id (27, 28, 29) # ringSize のユニーク値: ['0', '3', '4', '5', '6', '7', '8'] → 7種類 self.ringsize_embed = nn.Embedding(num_embeddings=7, embedding_dim=4) self.aroma_num_embed = nn.Embedding(num_embeddings=5, embedding_dim=4) # 0..4 self.fused_if_embed = nn.Embedding(num_embeddings=8, embedding_dim=4) # 0..7 # NEW: het27 (30) 用の embedding (0..26 → 27クラス) self.het27_embed = nn.Embedding(num_embeddings=27, embedding_dim=4) # 官能基フラグ 18 個 (各 2 次元) → 36 次元を 4 次元に圧縮 self.func_reduce = nn.Linear(18 * 2, 4) # ringSize の元の値 → index の対応（0,3,4,5,6,7,8 → 0..6） uniq = ['0', '3', '4', '5', '6', '7', '8'] uniq_int = sorted(int(x) for x in uniq) # [0,3,4,5,6,7,8] mapping = {v: i for i, v in enumerate(uniq_int)} self.register_buffer( "ring_values_tensor", torch.tensor(uniq_int, dtype=torch.long) ) self.ring_value_to_index = mapping # ---- bond_env_raw (31-78, 48 dims) 用の射影 ---- # 48 次元 → 16 次元に圧縮して concat self.bond_env_proj = nn.Linear(48, 16) # ---- 離散部分（x0..x6 + flags4 + x25,x26 + x27,x28,x29,x30）を 52→48 に圧縮 ---- # x0..x6: 7*4 = 28 # flags4: 4 # x25,x26: 2+2 = 4 # x27,x28,x29,x30: 4*4 = 16 # 合計 = 28 + 4 + 4 + 16 = 52 self.disc_proj = nn.Linear(52, 48) # このクラスの最終出力次元: # 離散部 (disc_proj 後) : 48 # bond_env_proj : 16 # 合計 = 64 self.out_dim = 48 + 16 @staticmethod def _build_element_lut(ELEMENTS): max_z = max(ELEMENTS) lut = torch.zeros(max_z + 1, dtype=torch.long) for i, z in enumerate(ELEMENTS): lut[z] = i return lut def forward(self, atom_inputs: torch.Tensor) -> torch.Tensor: """ Encode per-atom features into a dense embedding. Parameters ---------- atom_inputs : torch.Tensor Shape [N, 79]. 0 : Z 1 : degree 2 : charge 3 : hyb 4 : arom 5 : ring 6 : hcount 7-24 : func_flags (18) 25 : H-donor flag 26 : H-acceptor flag 27 : ringSize (0,3,4,5,6,7,8,...) 28 : #aromatic neighbors 29 : fused ring id (0..7) 30 : het27 (0..26) 31-78 : bond_env_raw (48) float Returns ------- out : torch.Tensor Shape [N, 64] = disc_emb(48) + bond_env_emb(16) """ device = next(self.parameters()).device x = atom_inputs.to(device=device, non_blocking=True) if x.ndim != 2 or x.size(1) != 79: raise ValueError(f"atom_inputs must be [N,79], got {tuple(x.shape)}") # ---------------------------- # helpers # ---------------------------- def clamp_idx(col: int, n_embed: int, lo: int = 0) -> torch.Tensor: """Read integer column and clamp to valid embedding range.""" return x[:, col].long().clamp(lo, n_embed - 1) def flag01(col: int) -> torch.Tensor: """Read 0/1 flag column.""" return x[:, col].long().clamp(0, 1) # ---------------------------- # core discrete features # ---------------------------- # 0: Z -> element_lut -> element_embed z_raw = x[:, 0].long() max_z = int(self.element_lut.shape[0] - 1) z_safe = torch.where((0 <= z_raw) & (z_raw <= max_z), z_raw, torch.zeros_like(z_raw)) idx0 = self.element_lut[z_safe].clamp(0, self.element_embed.num_embeddings - 1) x0 = self.element_embed(idx0) # 1: degree x1 = self.degree_embed(clamp_idx(1, self.degree_embed.num_embeddings)) # 2: ring flag (+1) for ring_embed (0..) ring_plus1 = (x[:, 5].long() + 1).clamp(0, self.ring_embed.num_embeddings - 1) x2 = self.ring_embed(ring_plus1) # 3: charge x3 = self.charge_embed(clamp_idx(2, self.charge_embed.num_embeddings)) # 4: aromatic (0/1) x4 = self.aromatic_embed(clamp_idx(4, self.aromatic_embed.num_embeddings)) # 5: hybrid x5 = self.hybrid_embed(clamp_idx(3, self.hybrid_embed.num_embeddings)) # 6: #hydrogens x6 = self.hydrogen_embed(clamp_idx(6, self.hydrogen_embed.num_embeddings)) # ---------------------------- # functional flags (7..24): 18 flags -> per-flag embed -> reduce # ---------------------------- # assumes each func_embed_i is an Embedding(2, d_flag) (you used 2 dims each before) func_embeds = [ self.func_embed_0, self.func_embed_1, self.func_embed_2, self.func_embed_3, self.func_embed_4, self.func_embed_5, self.func_embed_6, self.func_embed_7, self.func_embed_8, self.func_embed_9, self.func_embed_10, self.func_embed_11, self.func_embed_12, self.func_embed_13, self.func_embed_14, self.func_embed_15, self.func_embed_16, self.func_embed_17, ] flags = torch.cat([emb(flag01(7 + i)) for i, emb in enumerate(func_embeds)], dim=-1) # [N, 18*d_flag] flags4 = self.func_reduce(flags) # [N, 4] # 25,26: H-bond donor/acceptor flags x25 = self.h_don_embed(flag01(25)) x26 = self.h_acc_embed(flag01(26)) # ---------------------------- # ring size mapping (27) # ---------------------------- raw27 = x[:, 27].long() # default: last index = "other/unknown" mapped27 = torch.full_like(raw27, fill_value=self.ringsize_embed.num_embeddings - 1) # map known values via dict {value:int_index} # NOTE: this loop is small (few keys), so it's fine for v, idx in self.ring_value_to_index.items(): mapped27[raw27 == int(v)] = int(idx) mapped27 = mapped27.clamp(0, self.ringsize_embed.num_embeddings - 1) x27 = self.ringsize_embed(mapped27) # 28: #aromatic neighbors x28 = self.aroma_num_embed(clamp_idx(28, self.aroma_num_embed.num_embeddings)) # 29: fused ring id x29 = self.fused_if_embed(clamp_idx(29, self.fused_if_embed.num_embeddings)) # 30: het27 x30 = self.het27_embed(clamp_idx(30, self.het27_embed.num_embeddings)) # ---------------------------- # bond_env_raw (31..78): 48 floats -> proj -> 16 # ---------------------------- bond_env = x[:, 31:].to(dtype=torch.float32) if bond_env.size(1) != 48: raise ValueError(f"bond_env_raw must be 48 dims, got {bond_env.size(1)}") bond_env_emb = self.bond_env_proj(bond_env) # [N, 16] # ---------------------------- # concatenate discrete, project, then concat with bond env # ---------------------------- disc_cat = torch.cat( [ x0, x1, x2, x3, x4, x5, x6, # base discrete embeds flags4, # [N,4] x25, x26, # donor/acceptor x27, x28, x29, x30, # ringSize, aromNbrs, fusedId, het27 ], dim=-1, ) # expected [N, 52] in your original design disc_emb = self.disc_proj(disc_cat) # [N, 48] out = torch.cat([disc_emb, bond_env_emb], dim=-1) # [N, 64] return out class EquivariantThreeHopGINE(nn.Module): def __init__(self, in_feats, hidden_feats, out_feats, args): super().__init__() import torch if args is None: args = get_args() self.feat_embed = AtomEmbedding() # AtomEmbedding → 64 dims (48 from discrete + 16 from bond_env) self.linear_0 = nn.Linear(64, args.hidden_dim) # h0 edge_emb_dim = getattr(args, "edge_emb_dim", 32) self.bond_emb = nn.Embedding(5, edge_emb_dim, padding_idx=0) def mlp(in_f, out_f): return nn.Sequential( nn.Linear(in_f, out_f), nn.ReLU(), nn.Linear(out_f, out_f), nn.ReLU() ) nn1 = mlp(args.hidden_dim, hidden_feats) nn2 = mlp(hidden_feats, hidden_feats) nn3 = mlp(hidden_feats, hidden_feats) self.gine1 = GINEConv(nn1, edge_dim=edge_emb_dim) self.gine2 = GINEConv(nn2, edge_dim=edge_emb_dim) self.gine3 = GINEConv(nn3, edge_dim=edge_emb_dim) self.ln_in = nn.LayerNorm(args.hidden_dim) self.ln1 = nn.LayerNorm(hidden_feats) self.ln2 = nn.LayerNorm(hidden_feats) self.ln3 = nn.LayerNorm(hidden_feats) # Residual scales self.res1 = nn.Parameter(torch.tensor(0.5)) self.res2 = nn.Parameter(torch.tensor(0.5)) self.res3 = nn.Parameter(torch.tensor(0.5)) # If args.hidden_dim != hidden_feats, project skip for hop1 self.skip0 = None if args.hidden_dim != hidden_feats: self.skip0 = nn.Linear(args.hidden_dim, hidden_feats, bias=False) # JK: concat h0 + h1 + h2 + h3 jk_dim = args.hidden_dim + 3 * hidden_feats self.mix = nn.Sequential( nn.Linear(jk_dim, 2 * hidden_feats), nn.ReLU(), nn.Linear(2 * hidden_feats, hidden_feats), nn.ReLU(), ) # Project to VQ dim self.out_proj = nn.Linear(hidden_feats, args.hidden_dim) self.pre_vq_ln = nn.LayerNorm(args.hidden_dim) self.vq = VectorQuantize( dim=args.hidden_dim, codebook_size=args.codebook_size, decay=getattr(args, "ema_decay", 0.8), threshold_ema_dead_code=2, ) def reset_kmeans(self): self.vq._codebook.reset_kmeans() def forward(self, data, features, chunk_i, mask_dict=None, logger=None, epoch=None, batched_graph_base=None, mode=None, attr_list=None): import torch dev = next(self.parameters()).device if mode == "init_kmeans_final": if hasattr(data, "to"): data = data.to(dev) if torch.is_tensor(features): features = features.to(dev, non_blocking=True) self.vq(data, attr_list, mask_dict, logger, chunk_i, epoch, mode) return 0 # Edges (mirror for undirected) -> dev s1, d1 = data.edges() s1 = s1.to(dev, non_blocking=True); d1 = d1.to(dev, non_blocking=True) src = torch.cat([s1, d1], 0); dst = torch.cat([d1, s1], 0) edge_index = torch.stack([src, dst], 0) # Edge attributes (bond types {1..4}, 0 otherwise) eb = data.edata.get("weight", torch.zeros(data.num_edges(), dtype=torch.long, device=s1.device)) eb = eb.to(dev, non_blocking=True) e = torch.cat([eb, eb], 0) e = torch.where((e >= 1) & (e <= 4), e, torch.zeros_like(e)) edge_attr = self.bond_emb(e.long()) # Node features -> h0 features = features.to(dev, non_blocking=True) h0_raw = self.feat_embed(features) # [N, 64] h0 = self.ln_in(self.linear_0(h0_raw)) # [N, args.hidden_dim] # Hop 1 h0_for1 = self.skip0(h0) if self.skip0 is not None else h0 h1 = self.gine1(h0, edge_index, edge_attr) h1 = self.ln1(h1 * self.res1 + h0_for1) # Hop 2 h2 = self.gine2(h1, edge_index, edge_attr) h2 = self.ln2(h2 * self.res2 + h1) # Hop 3 h3 = self.gine3(h2, edge_index, edge_attr) h3 = self.ln3(h3 * self.res3 + h2) # JK concat h_cat = torch.cat([h0, h1, h2, h3], dim=-1) h_mid = self.mix(h_cat) h_out = self.out_proj(h_mid) # K-means 用のループ時は VQ を通さずにそのまま返す if mode == "init_kmeans_loop": return h_out import torch.nn.functional as F h_vq = self.pre_vq_ln(h_out) # L2 normalize latents here h_vq = F.normalize(h_vq, p=2, dim=-1, eps=1e-12) quantize_output = self.vq( h_vq, attr_list, mask_dict, logger, chunk_i, epoch, mode, ) if mode == "infer": # vq.infer returns: (kid_full, cid_full, gid_full, id2safe) if isinstance(quantize_output, (tuple, list)) and len(quantize_output) == 4: kid_full, cid_full, gid_full, id2safe = quantize_output # IMPORTANT: return format expected by evaluate/run_infer_after_restore # (they later do: kid, cid, gid, id2safe = ids) ids = (kid_full, cid_full, gid_full, id2safe) # The rest of your forward might expect (loss, ids, id2safe) or similar. # If you already return (_, ids, _) in evaluate, keep that structure. return None, ids, None # Backward compat (older vq): (gid, cid, id2safe) if isinstance(quantize_output, (tuple, list)) and len(quantize_output) == 3: gid_full, cid_full, id2safe = quantize_output kid_full = gid_full ids = (kid_full, cid_full, gid_full, id2safe) return None, ids, None raise TypeError(f"[model.forward] infer expected vq to return 4-tuple, got {type(quantize_output)}") # ------------------------- # INFER: return token IDs # ------------------------- if mode == "infer": # vq は infer のとき (key_id_full, cluster_id_full, id2safe) を返す設計にする # ここで normalize は使わない（loss群とは別物なので） if not (isinstance(quantize_output, (tuple, list)) and len(quantize_output) == 3): raise TypeError( f"[model.forward] mode='infer' expects vq() -> (key_id_full, cluster_id_full, id2safe), " f"got {type(quantize_output)} with len={len(quantize_output) if isinstance(quantize_output,(tuple,list)) else 'NA'}" ) key_id_full, cluster_id_full, id2safe = quantize_output # shape guarantee key_id_full = key_id_full.reshape(-1).long() cluster_id_full = cluster_id_full.reshape(-1).long() return key_id_full, cluster_id_full, id2safe # quantizer の生の出力（2タプルなど）を正規化して loss 群だけ取り出す loss, _embed_ignored, commit_loss, cb_loss, sil_loss, repel_loss, cb_repel_loss, ent_loss = _normalize_quantize_output( quantize_output, logger, device=h_vq.device if torch.is_tensor(h_vq) else None, dtype=getattr(h_vq, "dtype", None) ) # embed は今回は h_vq をそのまま使う（まだ本当の量子化はしていないバージョン） embed = h_vq # モデルが返すのは (total_loss, embed, loss_list) # return loss, embed, [commit_loss, repel_loss, cb_repel_loss] # loss, cb, loss_list3 = outputs return loss, embed, [commit_loss, cb_repel_loss, repel_loss, cb_loss, sil_loss, ent_loss] class Model(nn.Module): """ Wrapper of different models """ def __init__(self, conf): super(Model, self).__init__() self.model_name = conf["model_name"] if "MLP" in conf["model_name"]: self.encoder = MLP( num_layers=conf["num_layers"], input_dim=conf["feat_dim"], hidden_dim=conf["hidden_dim"], output_dim=conf["label_dim"], dropout_ratio=conf["dropout_ratio"], norm_type=conf["norm_type"], ).to(conf["device"]) elif "SAGE" in conf["model_name"]: self.encoder = SAGE( num_layers=conf["num_layers"], input_dim=conf["feat_dim"], hidden_dim=conf["hidden_dim"], output_dim=conf["label_dim"], dropout_ratio=conf["dropout_ratio"], activation=F.relu, norm_type=conf["norm_type"], codebook_size=conf["codebook_size"], lamb_edge=conf["lamb_edge"], lamb_node=conf["lamb_node"], lamb_div_ele=conf["lamb_div_ele"] ).to(conf["device"]) elif "GCN" in conf["model_name"]: self.encoder = GCN( num_layers=conf["num_layers"], input_dim=conf["feat_dim"], hidden_dim=conf["hidden_dim"], output_dim=conf["label_dim"], dropout_ratio=conf["dropout_ratio"], activation=F.relu, norm_type=conf["norm_type"], codebook_size=conf["codebook_size"], lamb_edge=conf["lamb_edge"], lamb_node=conf["lamb_node"] ).to(conf["device"]) elif "GAT" in conf["model_name"]: self.encoder = GAT( num_layers=conf["num_layers"], input_dim=conf["feat_dim"], hidden_dim=conf["hidden_dim"], output_dim=conf["label_dim"], dropout_ratio=conf["dropout_ratio"], activation=F.relu, attn_drop=conf["attn_dropout_ratio"], ).to(conf["device"]) elif "APPNP" in conf["model_name"]: self.encoder = APPNP( num_layers=conf["num_layers"], input_dim=conf["feat_dim"], hidden_dim=conf["hidden_dim"], output_dim=conf["label_dim"], dropout_ratio=conf["dropout_ratio"], activation=F.relu, norm_type=conf["norm_type"], ).to(conf["device"]) def forward(self, data, feats, epoch, logger): """ data: a graph `g` or a `dataloader` of blocks """ if "MLP" in self.model_name: return self.encoder(feats) else: return self.encoder(data, feats, epoch, logger) def forward_fitnet(self, data, feats): """ Return a tuple (h_list, h) h_list: intermediate hidden representation h: final output """ if "MLP" in self.model_name: return self.encoder(feats) else: return self.encoder(data, feats) def inference(self, data, feats): if "SAGE" in self.model_name: # return self.forward(data, feats) return self.encoder.inference(data, feats) else: return self.forward(data, feats)