feat: 修改损失值

docs/inpainting-generation-design.md

@@ -0,0 +1,75 @@
+# 基于矩形掩码修补的生成策略设计
+
+## 背景
+
+当前生成策略是从训练集采样地图，经 VQ-VAE 编码得到码本 z，再用**随机采样的 z**（而非真实 z）作为 MaskGIT 的条件进行三阶段级联生成。在当前数据规模下，模型难以从纯随机条件中学习到合理的生成分布，效果不佳。
+
+## 新策略目标
+
+降低任务难度：保持从训练集采样，但不再用随机 z，而是**在真实地图上施加随机矩形掩码**，让模型作为"修补（inpainting）"任务来完成。这样模型既有真实地图的结构先验（通过 z 条件），又有明确的局部填空目标。
+
+## 核心改动
+
+### 1. 掩码方式
+
+**仅使用矩形分块掩码**（`dataset.py` `std_mask` 的第二种模式），不采用散点随机掩码：
+
+- 在 13×13 网格上，反复放置随机矩形块（高宽 2~6），直到掩码格数达到目标比例
+- 掩码比例可调，默认沿用 Beta(2,2) 采样（范围 [5%, 100%]），同时支持指定固定比例用于对照实验
+
+### 2. 条件输入 z
+
+不再随机采样 z，改为使用**被掩码地图经 VQ-VAE 编码后得到的真实 z**。这意味着：
+
+- 编码器输入可能是部分掩码后的地图（掩码位置填 MASK_ID=7）
+- VQ-VAE 从可见部分推断整体结构，得到的 z 携带地图的全局语义信息
+- MaskGIT 以该 z 为条件，在掩码区域补全被遮盖的图块
+
+### 3. 生成流程
+
+仍然采用三阶段级联，但每个阶段统一使用同一份掩码布局：
+
+```
+输入: 完整地图 gt [13×13]
+  ↓
+生成矩形掩码 mask [13×13 bool]  (比率可调)
+  ↓
+inp = gt.copy(); inp[mask] = MASK_TOKEN
+  ↓
+VQ-VAE 编码 inp → z1, z2, z3
+  ↓
+Stage1: mg1(inp, z1) → 补全墙壁 (mask ∩ target=1)
+Stage2: mg2(inp1+stage1_out, z2) → 补全门/怪物/入口
+Stage3: mg3(inp2+stage2_out, z3) → 补全资源
+  ↓
+输出: 最终地图 + 逐阶段可视化对比
+```
+
+### 4. 与现有 `maskgit_sample` 的适配
+
+现有 `maskgit_sample` 已支持 `keep_fixed` 参数，天然适配修补场景：
+
+- `keep_fixed=True`（默认）：输入中已有的非掩码图块在生成过程中锁定不变，MaskGIT 仅填充掩码位
+- 需要新增：将矩形掩码 `mask` 传入函数，使模型只尝试修改掩码区域
+
+## 可调参数
+
+| 参数 | 说明 | 默认值 |
+|------|------|--------|
+| `mask_ratio` | 掩码比例，可指定固定值或 `None` 使用 Beta(2,2) 采样 | `None`（Beta 采样） |
+| `block_h_range` | 矩形高度范围 | `(2, 7)` |
+| `block_w_range` | 矩形宽度范围 | `(2, 7)` |
+| `keep_fixed` | 三阶段是否锁定已有结构 | `(True, True, True)` |
+
+## 实现要点
+
+1. 将 `std_mask` 的矩形掩码生成逻辑提取为独立函数 `rect_mask(ratio, h_range, w_range) -> np.ndarray`
+2. 新增 `inpaint_generate(map, mask, models, device)` 函数，包装完整修补流程
+3. 验证时对同一张地图使用不同 `mask_ratio`（如 0%，30%，50%，70%）生成多组结果，横向对比
+4. 可视化输出时同时展示：原始地图、掩码地图、各阶段生成结果、最终结果
+
+## 预期效果
+
+- 任务从"开放式生成"变为"约束式修补"，难度显著降低
+- 模型可利用可见区域的上下文信息（如墙面走向、房间布局）推断掩码区域内容
+- 不同掩码比例可衡量模型对空间结构的理解深度

ginka/dataset.py

@@ -10,10 +10,12 @@ def load_data(path: str):
         data = json.load(f)
 
     data_list = []
-    for value in data["data"].values():
+    map_names = []
+    for key, value in data["data"].items():
         data_list.append(value)
+        map_names.append(key)
 
-    return data_list
+    return data_list, map_names
 
 def compute_symmetry(target_np: np.ndarray) -> tuple:
     """从 numpy 地图矩阵中直接计算三种对称性，O(H*W)"""
@@ -39,7 +41,7 @@ class GinkaSeperatedDataset(Dataset):
         subset_weights: tuple = (0.5, 0.3, 0.2),
         density_stats: dict | None = None
     ):
-        self.data = load_data(data_path)
+        self.data, self.map_names = load_data(data_path)
         total = sum(subset_weights)
         self.subset_cumw = [sum(subset_weights[:i+1]) / total for i in range(len(subset_weights))]
 
@@ -130,6 +132,7 @@ class GinkaSeperatedDataset(Dataset):
             "distance_field": torch.LongTensor(compute_distance_field(enc1))
         }
         sample['sample_idx'] = idx
+        sample['map_name'] = self.map_names[idx]
         return sample
 
     def degrade_tile(self, m: np.ndarray, tiles: list) -> np.ndarray:
@@ -240,4 +243,6 @@ class GinkaSeperatedDataset(Dataset):
         else:
             out = self.apply_subset3(map_np)
 
-        return self.pack_sample(item, map_np, out)
+        sample = self.pack_sample(item, map_np, out)
+        sample["map_name"] = self.map_names[idx]
+        return sample

ginka/train_seperated.py

@@ -39,8 +39,8 @@ from shared.distance import DIST_VOCAB, compute_distance_field_tensor
 
 # 共用 VQ-VAE 超参
 # 三组编码器（vq1/vq2/vq3）共享相同超参，分别对三阶段地图上下文独立编码
-VQ_L = 16 # 码字序列长度（每个编码器输出 L 个码字，量化后合并为 L*3）
+VQ_L = 8 # 码字序列长度（每个编码器输出 L 个码字，量化后合并为 L*3）
-VQ_K = 32 # codebook 大小（离散码本条目数）
+VQ_K = 16 # codebook 大小（离散码本条目数）
 VQ_D_Z = 64 # 码字维度
 VQ_BETA = 1.0 # commit loss 权重（防止编码器输出漂离 codebook）
 VQ_GAMMA = 0.0 # entropy loss 权重（当前未启用）
@@ -52,9 +52,10 @@ VQ_NHEAD = 4 # VQ-VAE 多头注意力头数
 # 距离场编码器超参
 L_DIST = 4 # 距离场码字序列长度
 K_DIST = 16 # 距离场 codebook 大小
-DIST_D_MODEL = 128 # 距离场编码器模型维度
+DIST_D_MODEL = 256 # 距离场编码器模型维度
 DIST_LAYERS = 3 # 距离场编码器 Transformer 层数
-DIST_DIM_FF = 512 # 距离场编码器 FF 维度
+DIST_DIM_FF = 1024 # 距离场编码器 FF 维度
+DIST_NHEAD = 8 # 距离场编码器注意力头数
 VQ_BETA_DIST = 0.5 # 距离场 commit loss 权重
 
 # 第一阶段 MaskGIT 超参
@@ -104,6 +105,18 @@ RESOURCE_DENSITY_IDX = 4
 MG_Z_DROPOUT = 0.1 # z 隐变量 Dropout 概率
 MG_STRUCT_DROPOUT = 0.1 # 结构参量 Dropout 概率
 
+# 邻接损失权重（三阶段）
+LAMBDA_ADJ1 = 0.6
+LAMBDA_ADJ2 = 0.3
+LAMBDA_ADJ3 = 0.1
+
+# Patch 损失权重（三阶段）及核参数
+LAMBDA_PATCH1 = 0.5
+LAMBDA_PATCH2 = 0.5
+LAMBDA_PATCH3 = 0.5
+PATCH_KERNEL_SIZE = 5
+PATCH_SIGMA = 1.2
+
 # 损失参数
 VQ_BETA = 0.5 # 承诺损失权重
 
@@ -176,7 +189,7 @@ def build_model(device: torch.device):
     # 距离场编码器与量化器：将 L1 距离场编码为离散 latent z_dist
     dist_encoder = DistFieldEncoder(
         vocab=DIST_VOCAB, L=L_DIST, d_z=VQ_D_Z, d_model=DIST_D_MODEL,
-        nhead=VQ_NHEAD, num_layers=DIST_LAYERS, dim_ff=DIST_DIM_FF,
+        nhead=DIST_NHEAD, num_layers=DIST_LAYERS, dim_ff=DIST_DIM_FF,
         map_h=MAP_H, map_w=MAP_W
     ).to(device)
     dist_quantizer = VectorQuantizer(K=K_DIST, d_z=VQ_D_Z).to(device)
@@ -202,6 +215,72 @@ def cross_entropy_loss(logits, target):
     # logits: [B, L, C]，需转为 [B, C, L] 以匹配 cross_entropy 期望格式
     return F.cross_entropy(logits.permute(0, 2, 1), target)
 
+def adjacency_loss(logits, target):
+    # 邻接损失：约束相邻两格同时为空地的概率
+    # logits: [B, S, C] — MaskGIT 解码器输出
+    # target: [B, S] — 目标类别 ID，不含 MASK 标记
+    B, S, C = logits.shape
+    H = 13
+    W = 13
+    probs = F.softmax(logits, dim=-1) # [B, S, C]
+    p_floor = probs[:, :, 0].view(B, H, W) # [B, H, W] — 地板概率
+    t = target.view(B, H, W)
+    t_floor = (t == 0).float() # 地板标注为 1，其余为 0
+
+    # 水平边：左格 × 右格
+    joint_h = p_floor[:, :, :-1] * p_floor[:, :, 1:] # [B, H, W-1]
+    target_h = t_floor[:, :, :-1] * t_floor[:, :, 1:] # [B, H, W-1]
+
+    # 垂直边：上格 × 下格
+    joint_v = p_floor[:, :-1, :] * p_floor[:, 1:, :] # [B, H-1, W]
+    target_v = t_floor[:, :-1, :] * t_floor[:, 1:, :] # [B, H-1, W]
+
+    loss_h = F.binary_cross_entropy(joint_h, target_h, reduction='mean')
+    loss_v = F.binary_cross_entropy(joint_v, target_v, reduction='mean')
+    return (loss_h + loss_v) / 2.0
+
+def gaussian_kernel(kernel_size, sigma, device):
+    # 生成归一化二维高斯卷积核 [1, 1, K, K]
+    k = kernel_size
+    center = (k - 1) / 2.0
+    xs = torch.arange(k, dtype=torch.float32, device=device) - center
+    gx = torch.exp(-xs ** 2 / (2.0 * sigma ** 2))
+    gy = torch.exp(-xs ** 2 / (2.0 * sigma ** 2))
+    g2d = gx[:, None] * gy[None, :] # [K, K]
+    g2d = g2d / g2d.sum() # 归一化
+    return g2d.view(1, 1, k, k)
+
+def patch_loss(logits, target, kernel_size=5, sigma=1.2):
+    # Patch 损失：高斯核加权的邻域 CE 平滑损失
+    # logits: [B, S, C]
+    # target: [B, S]
+    B, S, C = logits.shape
+    H = 13
+    W = 13
+
+    # 逐格 CE（不做 reduction）
+    ce = F.cross_entropy(
+        logits.reshape(-1, C), target.reshape(-1), reduction='none'
+    ).view(B, H, W) # [B, H, W]
+
+    # 高斯核
+    kernel = gaussian_kernel(kernel_size, sigma, logits.device) # [1, 1, K, K]
+
+    # replicate 填充后用 unfold 提取邻域
+    pad = kernel_size // 2
+    ce_padded = F.pad(
+        ce.view(B, 1, H, W), (pad, pad, pad, pad), mode='replicate'
+    )
+    # patches: [B, K*K, H*W]
+    patches = F.unfold(ce_padded, kernel_size=(kernel_size, kernel_size))
+    patches = patches.view(B, kernel_size * kernel_size, H, W) # [B, K*K, H, W]
+
+    # 加权求和
+    k_flat = kernel.view(1, kernel_size * kernel_size, 1, 1)
+    smoothed = (patches * k_flat).sum(dim=1) # [B, H, W]
+
+    return smoothed.mean()
+
 def apply_z_dropout(
     z_q: torch.Tensor,
     mask_embedding: nn.Parameter,
@@ -541,7 +620,7 @@ def build_dataset_sample_case(
     }
 
 def sample_case_label(case: dict) -> str:
-    return f"train#{case['sample_idx']}"
+    return case["sample"]["map_name"]
 
 # 验证可视化 part1：3×3 网格；行1=编码器输入，行2=掩码输入，行3=三阶段预测（合并）
 def visualize_part1(batch, logits1, logits2, logits3, tile_dict):
@@ -573,7 +652,7 @@ def visualize_part1(batch, logits1, logits2, logits3, tile_dict):
     result3[inp3_np == MASK_TOKEN] = pred3[inp3_np == MASK_TOKEN]
 
     rows = [
-        [to_img(enc1_np), to_img(enc2_np), to_img(enc3_np)],
+        [annotate(to_img(enc1_np), batch["map_name"][0]), to_img(enc2_np), to_img(enc3_np)],
         [to_img(inp1_np), to_img(inp2_np), to_img(inp3_np)],
         [to_img(result1), to_img(result2), to_img(result3)],
     ]
@@ -614,7 +693,7 @@ def visualize_part2(batch, z_q, z_dist, models, device, tile_dict):
     target_density_cpu = batch["target_density"][0]
 
     rows = [
-        [to_img(enc1_np), to_img(enc2_np), to_img(enc3_np)],
+        [annotate(to_img(enc1_np), batch["map_name"][0]), to_img(enc2_np), to_img(enc3_np)],
         [
             annotate(to_img(inp1_np), kf_label),
             annotate_labels(to_img(auto_pred1_np), struct_cpu, target_density_cpu),
@@ -712,6 +791,12 @@ def validate(
     loss2_total = torch.Tensor([0]).to(device)
     loss3_total = torch.Tensor([0]).to(device)
     commit_total = torch.Tensor([0]).to(device)
+    adj1_total = torch.Tensor([0]).to(device)
+    adj2_total = torch.Tensor([0]).to(device)
+    adj3_total = torch.Tensor([0]).to(device)
+    patch1_total = torch.Tensor([0]).to(device)
+    patch2_total = torch.Tensor([0]).to(device)
+    patch3_total = torch.Tensor([0]).to(device)
     code_hits_total = torch.zeros(3, quantizer1.K, device=device)
 
     density_metrics = {
@@ -771,6 +856,12 @@ def validate(
             loss2_total += cross_entropy_loss(logits2, target2)
             loss3_total += cross_entropy_loss(logits3, target3)
             commit_total += commit_loss
+            adj1_total += adjacency_loss(logits1, target1)
+            adj2_total += adjacency_loss(logits2, target2)
+            adj3_total += adjacency_loss(logits3, target3)
+            patch1_total += patch_loss(logits1, target1, PATCH_KERNEL_SIZE, PATCH_SIGMA)
+            patch2_total += patch_loss(logits2, target2, PATCH_KERNEL_SIZE, PATCH_SIGMA)
+            patch3_total += patch_loss(logits3, target3, PATCH_KERNEL_SIZE, PATCH_SIGMA)
             code_hits_total += code_hits
 
             # 计算各目标对象的真实密度误差与过量生成密度
@@ -818,7 +909,7 @@ def validate(
     for m in [vq1, vq2, vq3, mg1, mg2, mg3, dist_encoder]:
         m.train()
 
-    return loss1_total, loss2_total, loss3_total, commit_total, code_hits_total
+    return loss1_total, loss2_total, loss3_total, adj1_total, adj2_total, adj3_total, patch1_total, patch2_total, patch3_total, commit_total, code_hits_total
 
 def train(device: torch.device):
     args = parse_arguments()
@@ -902,6 +993,12 @@ def train(device: torch.device):
         loss2_total = torch.Tensor([0]).to(device)
         loss3_total = torch.Tensor([0]).to(device)
         commit_total = torch.Tensor([0]).to(device)
+        adj1_total = torch.Tensor([0]).to(device)
+        adj2_total = torch.Tensor([0]).to(device)
+        adj3_total = torch.Tensor([0]).to(device)
+        patch1_total = torch.Tensor([0]).to(device)
+        patch2_total = torch.Tensor([0]).to(device)
+        patch3_total = torch.Tensor([0]).to(device)
         code_hits_total = torch.zeros(3, quantizer1.K, device=device)
 
         for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm):
@@ -966,15 +1063,25 @@ def train(device: torch.device):
             logits2 = mg2(inp2, z_q2, z_dist, struct, remain2)
             logits3 = mg3(inp3, z_q3, z_dist, struct, remain3)
 
-            # 三阶段 Cross Entropy + VQ commit loss 加权求和
+            # 三阶段 Cross Entropy + 邻接损失 + Patch 损失 + VQ commit loss 加权求和
             loss1 = cross_entropy_loss(logits1, target1)
             loss2 = cross_entropy_loss(logits2, target2)
             loss3 = cross_entropy_loss(logits3, target3)
+
+            adj1 = adjacency_loss(logits1, target1)
+            adj2 = adjacency_loss(logits2, target2)
+            adj3 = adjacency_loss(logits3, target3)
+            patch1 = patch_loss(logits1, target1, PATCH_KERNEL_SIZE, PATCH_SIGMA)
+            patch2 = patch_loss(logits2, target2, PATCH_KERNEL_SIZE, PATCH_SIGMA)
+            patch3 = patch_loss(logits3, target3, PATCH_KERNEL_SIZE, PATCH_SIGMA)
+
             loss1_weighted = STAGE1_CE_WEIGHT * loss1
             loss2_weighted = STAGE2_CE_WEIGHT * loss2
             loss3_weighted = STAGE3_CE_WEIGHT * loss3
+            adj_weighted = LAMBDA_ADJ1 * adj1 + LAMBDA_ADJ2 * adj2 + LAMBDA_ADJ3 * adj3
+            patch_weighted = LAMBDA_PATCH1 * patch1 + LAMBDA_PATCH2 * patch2 + LAMBDA_PATCH3 * patch3
             commit_weighted = VQ_BETA * commit_loss + VQ_BETA_DIST * commit_loss_dist
-            loss = loss1_weighted + loss2_weighted + loss3_weighted + commit_weighted
+            loss = loss1_weighted + loss2_weighted + loss3_weighted + adj_weighted + patch_weighted + commit_weighted
 
             loss.backward()
             optimizer.step()
@@ -985,6 +1092,12 @@ def train(device: torch.device):
             loss2_total += loss2.detach()
             loss3_total += loss3.detach()
             commit_total += commit_loss.detach()
+            adj1_total += adj1.detach()
+            adj2_total += adj2.detach()
+            adj3_total += adj3.detach()
+            patch1_total += patch1.detach()
+            patch2_total += patch2.detach()
+            patch3_total += patch3.detach()
             code_hits_total += code_hits.detach()
 
         # 每个 epoch 结束后更新学习率
@@ -994,11 +1107,12 @@ def train(device: torch.device):
         train_perplexity, train_usage_rate, train_active_codes = summarize_codebook_hits(code_hits_total)
         tqdm.write(
             f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
-            f"E: {epoch + 1} | Loss: {loss_total.item() / data_length:.6f} | "
+            f"E: {epoch + 1} | "
-            f"L1: {loss1_total.item() / data_length:.6f} | "
+            f"Loss: {loss_total.item() / data_length:.4f} | "
-            f"L2: {loss2_total.item() / data_length:.6f} | "
+            f"CE: {loss1_total.item() / data_length:.4f}, {loss2_total.item() / data_length:.4f}, {loss3_total.item() / data_length:.4f} | "
-            f"L3: {loss3_total.item() / data_length:.6f} | "
+            f"ADJ: {(LAMBDA_ADJ1 * adj1_total.item() + LAMBDA_ADJ2 * adj2_total.item() + LAMBDA_ADJ3 * adj3_total.item()) / data_length:.4f}, {adj1_total.item() / data_length:.4f}, {adj2_total.item() / data_length:.4f}, {adj3_total.item() / data_length:.4f} | "
-            f"VQ: {commit_total.item() / data_length:.6f} | "
+            f"PAT: {(LAMBDA_PATCH1 * patch1_total.item() + LAMBDA_PATCH2 * patch2_total.item() + LAMBDA_PATCH3 * patch3_total.item()) / data_length:.4f}, {patch1_total.item() / data_length:.4f}, {patch2_total.item() / data_length:.4f}, {patch3_total.item() / data_length:.4f} | "
+            f"VQ: {commit_total.item() / data_length:.4f} | "
             f"PPL: {train_perplexity:.4f} | "
             f"Usage: {train_usage_rate:.4f} ({train_active_codes}/{code_hits_total.numel()}) | "
             f"LR: {scheduler.get_last_lr()[0]:.6f}"
@@ -1009,22 +1123,25 @@ def train(device: torch.device):
             losses = validate(
                 dataloader_val, models, dist_models, device, tile_dict, dataset, epoch + 1
             )
-            loss1_total, loss2_total, loss3_total, commit_total, code_hits_total = losses
+            loss1_total, loss2_total, loss3_total, adj1_total, adj2_total, adj3_total, patch1_total, patch2_total, patch3_total, commit_total, code_hits_total = losses
             loss1_weighted = STAGE1_CE_WEIGHT * loss1_total
             loss2_weighted = STAGE2_CE_WEIGHT * loss2_total
             loss3_weighted = STAGE3_CE_WEIGHT * loss3_total
+            adj_weighted = LAMBDA_ADJ1 * adj1_total + LAMBDA_ADJ2 * adj2_total + LAMBDA_ADJ3 * adj3_total
+            patch_weighted = LAMBDA_PATCH1 * patch1_total + LAMBDA_PATCH2 * patch2_total + LAMBDA_PATCH3 * patch3_total
             commit_weighted = VQ_BETA * commit_total
-            loss_total = loss1_weighted + loss2_weighted + loss3_weighted + commit_weighted
+            loss_total = loss1_weighted + loss2_weighted + loss3_weighted + adj_weighted + patch_weighted + commit_weighted
 
             data_length = len(dataloader_val)
             val_perplexity, val_usage_rate, val_active_codes = summarize_codebook_hits(code_hits_total)
             tqdm.write(
                 f"[Validate {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
-                f"E: {epoch + 1} | Loss: {loss_total.item() / data_length:.6f} | "
+                f"E: {epoch + 1} | "
-                f"L1: {loss1_total.item() / data_length:.6f} | "
+                f"Loss: {loss_total.item() / data_length:.4f} | "
-                f"L2: {loss2_total.item() / data_length:.6f} | "
+                f"CE: {loss1_total.item() / data_length:.4f}, {loss2_total.item() / data_length:.4f}, {loss3_total.item() / data_length:.4f} | "
-                f"L3: {loss3_total.item() / data_length:.6f} | "
+                f"ADJ: {(LAMBDA_ADJ1 * adj1_total.item() + LAMBDA_ADJ2 * adj2_total.item() + LAMBDA_ADJ3 * adj3_total.item()) / data_length:.4f}, {adj1_total.item() / data_length:.4f}, {adj2_total.item() / data_length:.4f}, {adj3_total.item() / data_length:.4f} | "
-                f"VQ: {commit_total.item() / data_length:.6f} | "
+                f"PAT: {(LAMBDA_PATCH1 * patch1_total.item() + LAMBDA_PATCH2 * patch2_total.item() + LAMBDA_PATCH3 * patch3_total.item()) / data_length:.4f}, {patch1_total.item() / data_length:.4f}, {patch2_total.item() / data_length:.4f}, {patch3_total.item() / data_length:.4f} | "
+                f"VQ: {commit_total.item() / data_length:.4f} | "
                 f"PPL: {val_perplexity:.4f} | "
                 f"Usage: {val_usage_rate:.4f} ({val_active_codes}/{code_hits_total.numel()}) | "
             )