mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-08-14 18:12:28 +08:00
feat: 修改采样方式，改进码本
This commit is contained in:
parent
763f6258d2
commit
9aefa74714
@ -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 # 总训练轮数
|
||||
|
||||
238
ginka/sample.py
238
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():
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
if seed_mode:
|
||||
if stage1_method == "growth":
|
||||
pred1_np = wall_growth_sample(
|
||||
models.mg1, inp, z1, struct, target_density
|
||||
) # [B, H, W]
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
for img_i in range(5):
|
||||
samples_data = []
|
||||
for _ in range(4):
|
||||
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)
|
||||
|
||||
# 分层采样：将 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
|
||||
|
||||
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
|
||||
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)
|
||||
|
||||
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):
|
||||
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)
|
||||
ratio = random.uniform(0.2, 0.8)
|
||||
|
||||
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)
|
||||
|
||||
# 矩形分块掩码 + 确保墙壁连通性，与训练时一致
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
_, _, merged = full_generate(
|
||||
inp, z1, z2, z3, struct_t, target_density_t, models
|
||||
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)
|
||||
|
||||
samples_data.append(merged[0])
|
||||
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)
|
||||
|
||||
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)
|
||||
# 矩形分块掩码 + 确保墙壁连通性，与训练时一致
|
||||
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
|
||||
_, _, 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}")
|
||||
|
||||
Loading…
Reference in New Issue
Block a user