From 9aefa74714aa2b3dcebdb367776a476de44f8f96 Mon Sep 17 00:00:00 2001 From: unanmed <1319491857@qq.com> Date: Wed, 5 Aug 2026 00:28:32 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BF=AE=E6=94=B9=E9=87=87=E6=A0=B7?= =?UTF-8?q?=E6=96=B9=E5=BC=8F=EF=BC=8C=E6=94=B9=E8=BF=9B=E7=A0=81=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ginka/model.py | 36 ++--- ginka/sample.py | 242 +++++++++++++++++++++++------ ginka/train_seperated.py | 318 +++++++++++++++++++++++++-------------- 3 files changed, 418 insertions(+), 178 deletions(-) diff --git a/ginka/model.py b/ginka/model.py index a50175f..9db7a61 100644 --- a/ginka/model.py +++ b/ginka/model.py @@ -12,7 +12,7 @@ VQ_GAMMA = 0.1 # entropy loss 权重，鼓励码本使用均匀 # 三通道 VQ 各自独立超参（L、K、层数、维度等均独立配置） # Stage1 墙壁骨架 — 结构最复杂，模型容量最大 -VQ_L1 = 24 +VQ_L1 = 16 VQ_K1 = 32 VQ_D_MODEL1 = 384 VQ_NHEAD1 = 8 @@ -20,43 +20,43 @@ VQ_LAYERS1 = 6 VQ_DIM_FF1 = 1536 # Stage2 功能元素 — 中等复杂度 -VQ_L2 = 12 -VQ_K2 = 16 +VQ_L2 = 8 +VQ_K2 = 64 VQ_D_MODEL2 = 256 VQ_NHEAD2 = 4 -VQ_LAYERS2 = 6 +VQ_LAYERS2 = 4 VQ_DIM_FF2 = 1024 # Stage3 资源分布 — 最简单，模型容量最小 -VQ_L3 = 8 -VQ_K3 = 16 -VQ_D_MODEL3 = 192 -VQ_NHEAD3 = 4 +VQ_L3 = 24 +VQ_K3 = 24 +VQ_D_MODEL3 = 256 +VQ_NHEAD3 = 8 VQ_LAYERS3 = 4 -VQ_DIM_FF3 = 768 +VQ_DIM_FF3 = 1024 # 第一阶段 MaskGIT 超参 STAGE1_MG_DMODEL = 512 -STAGE1_MG_NHEAD = 4 +STAGE1_MG_NHEAD = 8 STAGE1_MG_NUM_LAYERS = 8 STAGE1_MG_DIM_FF = 2048 # 第二阶段 MaskGIT 超参 -STAGE2_MG_DMODEL = 256 -STAGE2_MG_NHEAD = 4 +STAGE2_MG_DMODEL = 384 +STAGE2_MG_NHEAD = 8 STAGE2_MG_NUM_LAYERS = 6 -STAGE2_MG_DIM_FF = 1024 +STAGE2_MG_DIM_FF = 1536 # 第三阶段 MaskGIT 超参 STAGE3_MG_DMODEL = 256 -STAGE3_MG_NHEAD = 4 +STAGE3_MG_NHEAD = 8 STAGE3_MG_NUM_LAYERS = 6 STAGE3_MG_DIM_FF = 1024 # 各阶段 VQ commit loss 权重（当前未单独使用，统一由 VQ_BETA 控制） -STAGE1_VQ_WEIGHT = 0.5 -STAGE2_VQ_WEIGHT = 0.5 -STAGE3_VQ_WEIGHT = 0.5 +STAGE1_VQ_WEIGHT = 0.2 +STAGE2_VQ_WEIGHT = 0.2 +STAGE3_VQ_WEIGHT = 0.2 # 全局参数 NUM_CLASSES = 8 # 图块类型数 @@ -65,7 +65,7 @@ TOTAL_K = VQ_K1 + VQ_K2 + VQ_K3 # 预计算，供日志输出使用 MAP_W = 13 # 地图宽度 MAP_H = 13 # 地图高度 -LR = 1e-4 # AdamW 初始学习率 +LR = 2e-4 # AdamW 初始学习率 MIN_LR = 1e-6 # 余弦退火最低学习率 WEIGHT_DECAY = 1e-4 # L2 正则化系数 EPOCHS = 400 # 总训练轮数 diff --git a/ginka/sample.py b/ginka/sample.py index 233df7a..12563f9 100644 --- a/ginka/sample.py +++ b/ginka/sample.py @@ -1,52 +1,96 @@ +import heapq import math +import random +from collections import deque import numpy as np import torch import torch.nn.functional as F from .model import MASK_TOKEN, MAP_H, MAP_W, SeperatedModels -from .utils import compute_remaining, MAP_SIZE, compute_adjacency_mask +from .utils import ( + compute_remaining, MAP_SIZE, compute_adjacency_mask, + WALL_DENSITY_IDX, DOOR_DENSITY_IDX, MONSTER_DENSITY_IDX, + ENTRANCE_DENSITY_IDX, RESOURCE_DENSITY_IDX +) -# MaskGIT 采样函数：通过迭代去掩码从离散隐变量 z 生成地图 +# 采样器超参 +DBC_GUMBEL = 2.0 # 揭开顺序的 Gumbel 噪声强度，0 关闭 +DBC_JITTER = 0.08 # 图块预算的随机浮动，0 关闭 +DBC_TRIES = 6 # 拒绝采样次数 + +# 每个阶段负责的图块，在 target_density 里对应的下标 +STAGE_DENSITY_IDX = { + 1: (WALL_DENSITY_IDX,), + 2: (DOOR_DENSITY_IDX, MONSTER_DENSITY_IDX, ENTRANCE_DENSITY_IDX), + 3: (RESOURCE_DENSITY_IDX,), +} + +_NB4 = ((0, 1), (0, -1), (1, 0), (-1, 0)) -def wall_growth_sample( - model: torch.nn.Module, - inp: torch.Tensor, - z: torch.Tensor, - struct: torch.Tensor, - target_density: torch.Tensor, - max_steps: int = 24 -) -> np.ndarray: - # 墙壁生长算法：从 inp 中已有的墙壁出发，逐步向外生长 - # 每步 MASK 位置决策（墙/非墙）后，新邻接面成为下一轮 MASK，逐步外扩 - state = inp.clone() # [B, MAP_SIZE] - # 初始 MASK：邻接已有墙壁的空地 - init_adj = compute_adjacency_mask(state) - state[init_adj & (state == 0)] = MASK_TOKEN +def floor_label(m: np.ndarray): + mask = (m != 1) + lab = -np.ones(m.shape, int) + n = 0 + for i in range(MAP_H): + for j in range(MAP_W): + if mask[i, j] and lab[i, j] < 0: + q = deque([(i, j)]) + lab[i, j] = n + while q: + a, b = q.popleft() + for da, db in _NB4: + x, y = a + da, b + db + if 0 <= x < MAP_H and 0 <= y < MAP_W and mask[x, y] and lab[x, y] < 0: + lab[x, y] = n + q.append((x, y)) + n += 1 + return lab, n - for step in range(max_steps): - mask_pos = (state == MASK_TOKEN) # [B, MAP_SIZE] - if not mask_pos.any(): - break - remain = compute_remaining(state, target_density, 1) - logits = model(state, z, struct, remain) - probs = F.softmax(logits, dim=-1) - - wall_prob = probs[:, :, 1] # [B, MAP_SIZE] - hits = (wall_prob > 0.5) & mask_pos - state[hits] = 1 - state[mask_pos & ~hits] = 0 - - adj = compute_adjacency_mask(state) - new_mask = adj & (state == 0) - if not new_mask.any(): - break - state[new_mask] = MASK_TOKEN - - state[state == MASK_TOKEN] = 0 - return state.cpu().numpy().reshape(state.size(0), MAP_H, MAP_W) +def repair_connectivity(m: np.ndarray) -> np.ndarray: + # 0-1 Dijkstra 打通最少的墙，使所有非墙格连成一片 + m = m.copy() + while True: + lab, n = floor_label(m) + if n <= 1: + return m + sizes = [int((lab == k).sum()) for k in range(n)] + main = int(np.argmax(sizes)) + tgt = next(k for k in range(n) if k != main) + dist = np.full(m.shape, 1 << 30) + pq = [] + prev = {} + end = None + for i in range(MAP_H): + for j in range(MAP_W): + if lab[i, j] == main: + dist[i, j] = 0 + heapq.heappush(pq, (0, i, j)) + while pq: + d, i, j = heapq.heappop(pq) + if d > dist[i, j]: + continue + if lab[i, j] == tgt: + end = (i, j) + break + for da, db in _NB4: + x, y = i + da, j + db + if not (0 <= x < MAP_H and 0 <= y < MAP_W): + continue + nd = d + (1 if m[x, y] == 1 else 0) + if nd < dist[x, y]: + dist[x, y] = nd + prev[(x, y)] = (i, j) + heapq.heappush(pq, (nd, x, y)) + if end is None: + return m + cur = end + while cur in prev: + if m[cur] == 1: + m[cur] = 0 + cur = prev[cur] def maskgit_sample( @@ -61,6 +105,13 @@ def maskgit_sample( current = inp.clone() target_tensor = torch.tensor(target_tiles, dtype=torch.long, device=inp.device) + # 退火对象改成「本阶段还需要放多少个目标图块」 + base = float(sum(target_density[0, i] for i in STAGE_DENSITY_IDX[stage])) * MAP_SIZE + if DBC_JITTER > 0: + base *= 1.0 + random.uniform(-DBC_JITTER, DBC_JITTER) + budget = int(round(base)) + need0 = max(0, budget - int(torch.isin(current[:], target_tensor).sum())) + # 迭代去掩码：每步根据置信度分数重新决定掩码位置 for step in range(steps): remain = compute_remaining(current, target_density, stage) @@ -72,20 +123,30 @@ def maskgit_sample( confidences = torch.gather(probs, -1, sampled.unsqueeze(-1)).squeeze(-1) - # 余弦退火调度：随步数推进，保留掩码的位置数量递减至 0 - ratio = math.cos(((step + 1) / steps) * math.pi / 2) - num_to_mask = math.floor(ratio * MAP_SIZE) + # Gumbel 噪声退火：提高揭开顺序的多样性（MaskGIT 原论文做法） + if DBC_GUMBEL > 0: + u = torch.rand_like(confidences).clamp_(1e-9, 1 - 1e-9) + score = (torch.log(confidences.clamp_min(1e-9)) + + DBC_GUMBEL * (1 - step / steps) * (-torch.log(-torch.log(u)))) + else: + score = confidences # 结构位：current 中非空地、非掩码的位置（来自上一阶段，始终保留） struct_mask = (current[:] != MASK_TOKEN) & (current[:] != 0) # 候选位：sampled 为目标图块且不覆盖结构位 candidate_mask = torch.isin(sampled[:], target_tensor) & ~struct_mask cand_count = candidate_mask.sum() - reveal_count = max(0, int(cand_count.item()) - num_to_mask) + + # 预算感知揭开：按本阶段还需放置的图块数做余弦退火 + ratio = math.cos(((step + 1) / steps) * math.pi / 2) + still = math.floor(ratio * need0) + now = max(0, budget - int(torch.isin(current[:], target_tensor).sum())) + reveal_count = min(max(0, now - still), int(cand_count.item())) + next_state = current[:].clone() if reveal_count > 0 and cand_count > 0: cand_indices = candidate_mask.nonzero(as_tuple=False) - cand_conf = confidences[:][cand_indices[:, 0], cand_indices[:, 1]] + cand_conf = score[:][cand_indices[:, 0], cand_indices[:, 1]] top_k = min(reveal_count, cand_conf.size(0)) _, top_idx = torch.topk(cand_conf, k=top_k, largest=True) reveal_rows = cand_indices[top_idx, 0] @@ -107,21 +168,108 @@ def maskgit_sample( return current.cpu().numpy().reshape(current.size(0), MAP_H, MAP_W) + +def sample_with_retry( + model: torch.nn.Module, inp: torch.Tensor, z: torch.Tensor, + struct: torch.Tensor, target_density: torch.Tensor, + stage: int, steps: int, target_tiles: list[int], + tries: int = DBC_TRIES +) -> np.ndarray: + # 拒绝采样：多次采样取地板连通的那张，都不连通则修复块数最少的 + best = None + best_n = 1 << 30 + for _ in range(tries): + m = maskgit_sample( + model, inp.clone(), z, struct, target_density, + stage, steps, target_tiles=target_tiles + ) + m_2d = m.reshape(MAP_H, MAP_W) + _, n = floor_label(m_2d) + if n == 1: + return m + if n < best_n: + best = m + best_n = n + if len(best.shape) == 3: + repaired = repair_connectivity(best[0]) + repaired = repaired[np.newaxis, ...] + else: + repaired = repair_connectivity(best.reshape(MAP_H, MAP_W)) + repaired = repaired.reshape(best.shape) + return repaired + + +def wall_growth_sample( + model: torch.nn.Module, + inp: torch.Tensor, + z: torch.Tensor, + struct: torch.Tensor, + target_density: torch.Tensor, + max_steps: int = 24 +) -> np.ndarray: + # 墙壁生长算法：从 inp 中已有的墙壁出发，逐步向外生长 + # 所有未固定位置初始为 MASK，只有与墙壁邻接的 MASK 才参与评估 + state = inp.clone() # [B, MAP_SIZE] + state[state == 0] = MASK_TOKEN + + for step in range(max_steps): + adj = compute_adjacency_mask(state) + mask_pos = adj & (state == MASK_TOKEN) + if not mask_pos.any(): + break + + remain = compute_remaining(state, target_density, 1) + logits = model(state, z, struct, remain) + probs = F.softmax(logits, dim=-1) + + wall_prob = probs[:, :, 1] # [B, MAP_SIZE] + hits = (wall_prob > 0.5) & mask_pos + state[hits] = 1 + state[mask_pos & ~hits] = 0 + + state[state == MASK_TOKEN] = 0 + return state.cpu().numpy().reshape(state.size(0), MAP_H, MAP_W) + + def full_generate( inp: torch.Tensor, z1: torch.Tensor, z2: torch.Tensor, z3: torch.Tensor, struct: torch.Tensor, target_density: torch.Tensor, models: SeperatedModels, - steps: int = 18 + steps: int = 18, + seed_mode: bool = False, + stage1_method: str = "maskgit" ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - # 三阶段级联生成：Stage1 墙壁生长 → Stage2 门/怪/入口 → Stage3 资源 + # 三阶段级联生成 + # seed_mode=True 从空白生成墙壁（自由生成） + # seed_mode=False 从 inp 中的已有结构出发（掩码补全） + # stage1_method: "maskgit"=MaskGIT+拒绝采样, "growth"=墙壁生长算法 # 返回 (stage1结果, stage1+2合并, 最终完整地图)，形状均为 [B, H, W] device = inp.device - pred1_np = wall_growth_sample( - models.mg1, inp, z1, struct, target_density - ) # [B, H, W] + if seed_mode: + if stage1_method == "growth": + pred1_np = wall_growth_sample( + models.mg1, inp, z1, struct, target_density + ) + else: + stage1_inp = torch.full((inp.size(0), MAP_SIZE), MASK_TOKEN, dtype=torch.long, device=device) + pred1_np = sample_with_retry( + models.mg1, stage1_inp, z1, struct, target_density, 1, + steps, target_tiles=[1] + ) + else: + if stage1_method == "growth": + pred1_np = wall_growth_sample( + models.mg1, inp, z1, struct, target_density + ) + else: + pred1_np = sample_with_retry( + models.mg1, inp, z1, struct, target_density, 1, + steps, target_tiles=[1] + ) + # [B, H, W] inp2 = torch.tensor(pred1_np.reshape(pred1_np.shape[0], -1), dtype=torch.long, device=device) inp2[inp2 == 0] = MASK_TOKEN diff --git a/ginka/train_seperated.py b/ginka/train_seperated.py index 96e2b37..50e868b 100644 --- a/ginka/train_seperated.py +++ b/ginka/train_seperated.py @@ -20,8 +20,8 @@ from .utils import ( compute_remaining, MAP_SIZE, summarize_codebook_hits ) from .sample import full_generate -from .dataset import GinkaSeperatedDataset -from shared.image import matrix_to_image_cv +from .dataset import GinkaSeperatedDataset, rect_mask, ensure_wall_connection +from shared.image import matrix_to_image_cv, annotate # 图块 ID 定义： # 0. 空地 1. 墙壁 2. 普通门 3. 资源 4. 怪物 5. 入口 6. 机关门 7. 掩码（MASK_TOKEN） @@ -107,133 +107,216 @@ def quantize_stage_latents( code_hits = (code_hits1, code_hits2, code_hits3) return (z_q1, z_q2, z_q3), commit_loss, code_hits, entropy_loss -# 墙壁种子生成：随机放置墙壁后让模型生长 -def visualize_seed( +# 每张图 2 行 x 4 列：MASKED | PREDICTED | MASKED | PREDICTED +def make_compare_grid(samples_src, samples_pred, map_keys, tile_dict, TILE_SIZE): + SEP = 3 + img_h = MAP_H * TILE_SIZE + img_w = MAP_W * TILE_SIZE + cols = 4 + rows = 2 + grid = np.ones((rows * img_h + (rows + 1) * SEP, cols * img_w + (cols + 1) * SEP, 3), dtype=np.uint8) * 255 + for r in range(rows): + row_y = SEP + r * (img_h + SEP) + for c in range(cols): + is_src = (c % 2 == 0) + idx = r * 2 + c // 2 + col_x = SEP + c * (img_w + SEP) + data = samples_src[idx] if is_src else samples_pred[idx] + grid[row_y:row_y + img_h, col_x:col_x + img_w] = annotate( + matrix_to_image_cv(data, tile_dict, TILE_SIZE), map_keys[idx], y=14 + ) + # 每两张图片之间的分隔线 + for c in range(cols - 1): + line_x = (c + 1) * (img_w + SEP) + SEP // 2 + grid[:, line_x:line_x + SEP, :] = 180 + return grid + +def visualize_seed_growth( train_dataset: GinkaSeperatedDataset, models: SeperatedModels, device: torch.device, tile_dict, epoch: int ): - SEP = 3 - TILE_SIZE = 32 - img_h = MAP_H * TILE_SIZE - img_w = MAP_W * TILE_SIZE - - def to_img(mat): - return matrix_to_image_cv(mat, tile_dict, TILE_SIZE) - save_dir = f"result/seperated/e{epoch}" os.makedirs(save_dir, exist_ok=True) + samples_src = [] + samples_pred = [] + map_keys = [] + for _ in range(20): + sample = train_dataset.random_sample_map() + struct_t = sample["struct_inject"].to(device).reshape(1, -1) + target_density_t = sample["target_density"].to(device).reshape(1, -1) + enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE) + enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE) + enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE) + enc1_np = enc1_t.cpu().numpy().reshape(MAP_H, MAP_W) + with torch.no_grad(): + z_e1 = models.vq1(enc1_t) + z_e2 = models.vq2(enc2_t) + z_e3 = models.vq3(enc3_t) + z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3) + z1, z2, z3 = z_q + inp = torch.full((1, MAP_SIZE), 0, dtype=torch.long, device=device) + seed_count = random.randint(7, 14) + seed_idx = torch.randperm(MAP_SIZE, device=device)[:seed_count] + inp[0, seed_idx] = 1 + inp_vis = inp.clone() + inp_vis[inp_vis == 0] = MASK_TOKEN + _, _, merged = full_generate( + inp, z1, z2, z3, struct_t, target_density_t, models, + steps=SEED_SAMPLE_STEPS, seed_mode=True, stage1_method="growth" + ) + samples_src.append(inp_vis.cpu().numpy().reshape(MAP_H, MAP_W)) + samples_pred.append(merged[0]) + map_keys.append(sample["map_name"]) + for i in range(5): + grid = make_compare_grid( + samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4], + map_keys[i*4:i*4+4], tile_dict, 32 + ) + cv2.imwrite(f"{save_dir}/seed_growth_{i}.png", grid) - for img_i in range(5): - samples_data = [] - for _ in range(4): - sample = train_dataset.random_sample_map() - struct_t = sample["struct_inject"].to(device).reshape(1, -1) - target_density_t = sample["target_density"].to(device).reshape(1, -1) - enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE) - enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE) - enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE) - - with torch.no_grad(): - z_e1 = models.vq1(enc1_t) - z_e2 = models.vq2(enc2_t) - z_e3 = models.vq3(enc3_t) - z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3) - z1, z2, z3 = z_q - - inp = torch.full((1, MAP_SIZE), 0, dtype=torch.long, device=device) - seed_count = random.randint(7, 14) - - # 分层采样：将 13x13 地图划分为 4x4 网格，从不同格中随机取种子点 - grid_h, grid_w = 4, 4 - cell_h = MAP_H // grid_h + 1 - cell_w = MAP_W // grid_w + 1 - all_cells = [(r, c) for r in range(grid_h) for c in range(grid_w)] - chosen_cells = random.sample(all_cells, seed_count) - for cell_r, cell_c in chosen_cells: - r_min = cell_r * cell_h - r_max = min(r_min + cell_h, MAP_H) - c_min = cell_c * cell_w - c_max = min(c_min + cell_w, MAP_W) - sr = random.randint(r_min, max(r_min, r_max - 1)) - sc = random.randint(c_min, max(c_min, c_max - 1)) - flat_idx = sr * MAP_W + sc - inp[0, flat_idx] = 1 - - _, _, merged = full_generate( - inp, z1, z2, z3, - struct_t, target_density_t, models, - steps=SEED_SAMPLE_STEPS - ) - - samples_data.append(merged[0]) - - grid = np.ones((2 * img_h + 3 * SEP, 2 * img_w + 3 * SEP, 3), dtype=np.uint8) * 255 - for r in range(2): - for c in range(2): - y = SEP + r * (img_h + SEP) - x = SEP + c * (img_w + SEP) - grid[y:y + img_h, x:x + img_w] = to_img(samples_data[r * 2 + c]) - cv2.imwrite(f"{save_dir}/seed{img_i}.png", grid) - -# 矩形掩码生成：对真实地图随机掩码后让模型补全 -def visualize_mask( +def visualize_seed_maskgit( train_dataset: GinkaSeperatedDataset, models: SeperatedModels, device: torch.device, tile_dict, epoch: int ): - SEP = 3 - TILE_SIZE = 32 - img_h = MAP_H * TILE_SIZE - img_w = MAP_W * TILE_SIZE - - def to_img(mat): - return matrix_to_image_cv(mat, tile_dict, TILE_SIZE) - save_dir = f"result/seperated/e{epoch}" os.makedirs(save_dir, exist_ok=True) + samples_src = [] + samples_pred = [] + map_keys = [] + for _ in range(20): + sample = train_dataset.random_sample_map() + struct_t = sample["struct_inject"].to(device).reshape(1, -1) + target_density_t = sample["target_density"].to(device).reshape(1, -1) + enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE) + enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE) + enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE) + enc1_np = enc1_t.cpu().numpy().reshape(MAP_H, MAP_W) + with torch.no_grad(): + z_e1 = models.vq1(enc1_t) + z_e2 = models.vq2(enc2_t) + z_e3 = models.vq3(enc3_t) + z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3) + z1, z2, z3 = z_q + inp = torch.full((1, MAP_SIZE), MASK_TOKEN, dtype=torch.long, device=device) + _, _, merged = full_generate( + inp, z1, z2, z3, struct_t, target_density_t, models, + steps=GENERATE_STEP, seed_mode=True, stage1_method="maskgit" + ) + samples_src.append(inp.cpu().numpy().reshape(MAP_H, MAP_W)) + samples_pred.append(merged[0]) + map_keys.append(sample["map_name"]) + for i in range(5): + grid = make_compare_grid( + samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4], + map_keys[i*4:i*4+4], tile_dict, 32 + ) + cv2.imwrite(f"{save_dir}/seed_maskgit_{i}.png", grid) - for img_i in range(5): - samples_data = [] - for _ in range(4): - sample = train_dataset.random_sample_map() - raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W) - ratio = random.uniform(0.2, 0.8) +def visualize_mask_growth( + train_dataset: GinkaSeperatedDataset, + models: SeperatedModels, + device: torch.device, + tile_dict, + epoch: int +): + save_dir = f"result/seperated/e{epoch}" + os.makedirs(save_dir, exist_ok=True) + samples_src = [] + samples_pred = [] + map_keys = [] + for _ in range(20): + sample = train_dataset.random_sample_map() + raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W) + enc1_np = sample["encoder_stage1"].cpu().numpy().reshape(MAP_H, MAP_W) + enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE) + enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE) + enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE) + struct_t = sample["struct_inject"].to(device).reshape(1, -1) + target_density_t = sample["target_density"].to(device).reshape(1, -1) - enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE) - enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE) - enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE) - struct_t = sample["struct_inject"].to(device).reshape(1, -1) - target_density_t = sample["target_density"].to(device).reshape(1, -1) + # 矩形分块掩码 + 确保墙壁连通性，与训练时一致 + ratio = random.uniform(0.3, 0.9) + rmask = rect_mask(ratio).reshape(MAP_H, MAP_W) + rmask = ensure_wall_connection(rmask, enc1_np) + inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE) + inp[0, rmask.flatten()] = MASK_TOKEN + masked_np = inp.cpu().numpy().reshape(MAP_H, MAP_W) - with torch.no_grad(): - z_e1 = models.vq1(enc1_t) - z_e2 = models.vq2(enc2_t) - z_e3 = models.vq3(enc3_t) - z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3) - z1, z2, z3 = z_q + with torch.no_grad(): + z_e1 = models.vq1(enc1_t) + z_e2 = models.vq2(enc2_t) + z_e3 = models.vq3(enc3_t) + z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3) + z1, z2, z3 = z_q + _, _, merged = full_generate( + inp, z1, z2, z3, struct_t, target_density_t, models, + stage1_method="growth" + ) + samples_src.append(masked_np) + samples_pred.append(merged[0]) + map_keys.append(sample["map_name"]) + for i in range(5): + grid = make_compare_grid( + samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4], + map_keys[i*4:i*4+4], tile_dict, 32 + ) + cv2.imwrite(f"{save_dir}/mask_growth_{i}.png", grid) - mask = torch.rand(MAP_SIZE, device=device) < ratio - inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE) - inp[0, mask] = MASK_TOKEN +def visualize_mask_maskgit( + train_dataset: GinkaSeperatedDataset, + models: SeperatedModels, + device: torch.device, + tile_dict, + epoch: int +): + save_dir = f"result/seperated/e{epoch}" + os.makedirs(save_dir, exist_ok=True) + samples_src = [] + samples_pred = [] + map_keys = [] + for _ in range(20): + sample = train_dataset.random_sample_map() + raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W) + enc1_np = sample["encoder_stage1"].cpu().numpy().reshape(MAP_H, MAP_W) + enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE) + enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE) + enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE) + struct_t = sample["struct_inject"].to(device).reshape(1, -1) + target_density_t = sample["target_density"].to(device).reshape(1, -1) - _, _, merged = full_generate( - inp, z1, z2, z3, struct_t, target_density_t, models - ) + # 矩形分块掩码 + 确保墙壁连通性，与训练时一致 + ratio = random.uniform(0.3, 0.9) + rmask = rect_mask(ratio).reshape(MAP_H, MAP_W) + rmask = ensure_wall_connection(rmask, enc1_np) + inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE) + inp[0, rmask.flatten()] = MASK_TOKEN + masked_np = inp.cpu().numpy().reshape(MAP_H, MAP_W) - samples_data.append(merged[0]) - - grid = np.ones((2 * img_h + 3 * SEP, 2 * img_w + 3 * SEP, 3), dtype=np.uint8) * 255 - for r in range(2): - for c in range(2): - y = SEP + r * (img_h + SEP) - x = SEP + c * (img_w + SEP) - grid[y:y + img_h, x:x + img_w] = to_img(samples_data[r * 2 + c]) - cv2.imwrite(f"{save_dir}/mask{img_i}.png", grid) + with torch.no_grad(): + z_e1 = models.vq1(enc1_t) + z_e2 = models.vq2(enc2_t) + z_e3 = models.vq3(enc3_t) + z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3) + z1, z2, z3 = z_q + _, _, merged = full_generate( + inp, z1, z2, z3, struct_t, target_density_t, models, + stage1_method="maskgit" + ) + samples_src.append(masked_np) + samples_pred.append(merged[0]) + map_keys.append(sample["map_name"]) + for i in range(5): + grid = make_compare_grid( + samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4], + map_keys[i*4:i*4+4], tile_dict, 32 + ) + cv2.imwrite(f"{save_dir}/mask_maskgit_{i}.png", grid) def train(device: torch.device): args = parse_arguments() @@ -283,6 +366,7 @@ 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) + entropy_total = torch.Tensor([0]).to(device) code_hits_total = (torch.zeros(result.quantizer1.K, device=device), torch.zeros(result.quantizer2.K, device=device), torch.zeros(result.quantizer3.K, device=device)) # validate for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm): @@ -330,10 +414,12 @@ def train(device: torch.device): logits2 = result.mg2(inp2, z_q2, struct, remain2) logits3 = result.mg3(inp3, z_q3, struct, remain3) - # 三阶段 Cross Entropy：仅对输入中为 MASK_TOKEN 的位置计算 loss + # 三阶段 Cross Entropy + # Stage1: 仅对掩码位置计算 loss + # Stage2/3: 掩码 + 空地位置均计算 loss，模型需学习哪里不能放东西 mask1 = (inp1 == MASK_TOKEN) - mask2 = (inp2 == MASK_TOKEN) - mask3 = (inp3 == MASK_TOKEN) + mask2 = (inp2 == MASK_TOKEN) | (inp2 == 0) + mask3 = (inp3 == MASK_TOKEN) | (inp3 == 0) loss1 = cross_entropy_loss(logits1, target1, mask1) loss2 = cross_entropy_loss(logits2, target2, mask2) loss3 = cross_entropy_loss(logits3, target3, mask3) @@ -353,12 +439,16 @@ def train(device: torch.device): loss2_total += loss2.detach() loss3_total += loss3.detach() commit_total += commit_loss.detach() + entropy_total += entropy_loss.detach() code_hits_total = (code_hits_total[0] + code_hits[0].detach(), code_hits_total[1] + code_hits[1].detach(), code_hits_total[2] + code_hits[2].detach()) # accumulate train # 每个 epoch 结束后更新学习率 result.scheduler.step() data_length = len(dataloader) + # 总 loss 去掉 entropy_loss 以便与 CE + VQ 对齐 + loss_display = (loss_total.item() - VQ_GAMMA * entropy_total.item()) / data_length + entropy_display = entropy_total.item() / data_length stats = summarize_codebook_hits(code_hits_total) parts = [] for name in ["q1(stage1)", "q2(stage2)", "q3(stage3)"]: @@ -368,9 +458,9 @@ def train(device: torch.device): tqdm.write( f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " f"E: {epoch + 1} | " - f"Loss: {loss_total.item() / data_length:.4f} | " + f"Loss: {loss_display:.4f} | " f"CE: {loss1_total.item() / data_length:.4f}, {loss2_total.item() / data_length:.4f}, {loss3_total.item() / data_length:.4f} | " - f"VQ: {commit_total.item() / data_length:.4f} | " + f"cmt: {commit_total.item() / data_length:.4f} ent: {entropy_display:+.4f} | " f"VQ: {' | '.join(parts)} | " f"Total: {c['active']}/{c['K']} ppl={c['ppl']:.1f} | " f"LR: {result.scheduler.get_last_lr()[0]:.6f}" @@ -378,8 +468,10 @@ def train(device: torch.device): # 每 CHECKPOINT 个 epoch 执行可视化并保存检查点 if (epoch + 1) % CHECKPOINT == 0: - visualize_seed(dataset, result, device, tile_dict, epoch + 1) - visualize_mask(dataset, result, device, tile_dict, epoch + 1) + visualize_seed_growth(dataset, result, device, tile_dict, epoch + 1) + visualize_seed_maskgit(dataset, result, device, tile_dict, epoch + 1) + visualize_mask_growth(dataset, result, device, tile_dict, epoch + 1) + visualize_mask_maskgit(dataset, result, device, tile_dict, epoch + 1) ckpt_path = f"result/seperated/sep-{epoch + 1}.pth" result.save(ckpt_path, epoch + 1) tqdm.write(f"Saved checkpoint: {ckpt_path}")