| """DeepSets 编码器 + LSTM 解码器（方案 B 自包含分发包的平铺依赖副本）。""" |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from vocab import Vocab |
| import config |
|
|
|
|
| class DeepSetsEncoder(nn.Module): |
| def __init__(self, input_dim: int = None): |
| super().__init__() |
| dims = config.ENCODER_MLP_DIMS |
| H = config.ENCODER_OUT_DIM |
|
|
| layers = [] |
| in_dim = input_dim if input_dim is not None else 0 |
| for out_dim in dims: |
| layers.append(nn.Linear(in_dim, out_dim)) |
| layers.append(nn.ReLU()) |
| in_dim = out_dim |
| self.phi = nn.Sequential(*layers) |
|
|
| rho_layers = [ |
| nn.Linear(H * 2, H), |
| nn.ReLU(), |
| nn.Linear(H, H), |
| ] |
| self.rho = nn.Sequential(*rho_layers) |
|
|
| def forward(self, X, y, n_samples, var_counts=None): |
| """ |
| X: [B, N_max, D_MAX] (padded, 0 for missing samples and missing vars) |
| y: [B, N_max] |
| n_samples: [B] 每个 batch 的真实样本数 |
| var_counts: [B] 每个样本的真实变量维度（未使用，保留兼容签名） |
| """ |
| B, N_max, D = X.shape |
| v = torch.cat([X, y.unsqueeze(-1)], dim=-1) |
|
|
| h = self.phi(v) |
|
|
| |
| sample_mask = torch.arange(N_max, device=X.device).unsqueeze(0) < n_samples.unsqueeze(1) |
| sample_mask = sample_mask.unsqueeze(-1).float() |
|
|
| h = h * sample_mask |
|
|
| h_mean = h.sum(dim=1) / n_samples.unsqueeze(-1).float().clamp(min=1) |
| h_max = h.masked_fill(sample_mask == 0, float("-inf")).max(dim=1).values |
|
|
| z = torch.cat([h_mean, h_max], dim=-1) |
| z = self.rho(z) |
| return z |
|
|
|
|
| class ExpressionLSTM(nn.Module): |
| def __init__(self, vocab_size: int): |
| super().__init__() |
| self.embedding = nn.Embedding(vocab_size, config.LSTM_EMBEDDING_DIM) |
| self.lstm = nn.LSTM( |
| input_size=config.LSTM_EMBEDDING_DIM, |
| hidden_size=config.LSTM_HIDDEN_DIM, |
| num_layers=config.LSTM_NUM_LAYERS, |
| batch_first=True, |
| dropout=config.DROPOUT if config.LSTM_NUM_LAYERS > 1 else 0, |
| ) |
| self.proj = nn.Linear(config.LSTM_HIDDEN_DIM, vocab_size) |
|
|
| def forward(self, token_ids, hidden, cell): |
| """ |
| token_ids: [B, T] |
| hidden, cell: [1, B, H] |
| returns: logits [B, T, V], hidden, cell |
| """ |
| embedded = self.embedding(token_ids) |
| output, (hidden, cell) = self.lstm(embedded, (hidden, cell)) |
| logits = self.proj(output) |
| return logits, hidden, cell |
|
|
|
|
| class SymbolicRegressionModel(nn.Module): |
| def __init__(self, vocab: Vocab): |
| super().__init__() |
| self.vocab = vocab |
| self.encoder = DeepSetsEncoder(input_dim=config.D_MAX + 1) |
| self.lstm = ExpressionLSTM(vocab.vocab_size) |
|
|
| def forward(self, batch): |
| X = batch["X"] |
| y = batch["y"] |
| n_samples = batch["n_samples"] |
| token_ids = batch["token_ids"] |
| var_counts = batch["var_counts"] |
|
|
| z = self.encoder(X, y, n_samples, var_counts) |
|
|
| hidden = z.unsqueeze(0) |
| cell = z.unsqueeze(0) |
|
|
| logits, _, _ = self.lstm(token_ids[:, :-1], hidden, cell) |
| return logits |
|
|
| @torch.no_grad() |
| def generate(self, X, y, n_samples, var_counts=None, max_len: int = None): |
| self.eval() |
| if max_len is None: |
| max_len = config.T_MAX |
|
|
| B = X.size(0) |
| if var_counts is None: |
| var_counts = torch.full((B,), X.size(-1), device=X.device) |
|
|
| z = self.encoder(X, y, n_samples, var_counts) |
|
|
| hidden = z.unsqueeze(0) |
| cell = z.unsqueeze(0) |
|
|
| token = torch.full((B, 1), fill_value=self.vocab.sos_id, device=X.device) |
| generated = [token] |
|
|
| for _ in range(max_len): |
| embedded = self.lstm.embedding(token) |
| output, (hidden, cell) = self.lstm.lstm(embedded, (hidden, cell)) |
| logits = self.lstm.proj(output) |
| token = logits.argmax(dim=-1) |
|
|
| generated.append(token) |
|
|
| if (token == self.vocab.eos_id).all(): |
| break |
|
|
| all_tokens = torch.cat(generated, dim=1) |
| return all_tokens |
|
|
