mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-05-14 04:41:12 +08:00
271 lines
11 KiB
Python
271 lines
11 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
import random
|
|
from datetime import datetime
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.optim as optim
|
|
import cv2
|
|
import numpy as np
|
|
from torch_geometric.loader import DataLoader
|
|
from tqdm import tqdm
|
|
from .transformer_vae.vae import GinkaTransformerVAE
|
|
from .vae_rnn.loss import VAELoss
|
|
from .vae_rnn.scheduler import VAEScheduler
|
|
from .dataset import GinkaRNNDataset
|
|
from shared.image import matrix_to_image_cv
|
|
|
|
# 手工标注标签定义（暂时不用）：
|
|
# 0. 蓝海, 1. 红海, 2: 室内, 3. 野外, 4. 左右对称, 5. 上下对称, 6. 伪对称, 7. 咸鱼层,
|
|
# 8. 剧情层, 9. 水层, 10. 爽塔, 11. Boss层, 12. 纯Boss层, 13. 多房间, 14. 多走廊, 15. 道具风
|
|
# 16. 区域入口, 17. 区域连接, 18. 有机关门, 19. 道具层, 20. 斜向对称, 21. 左右通道, 22. 上下通道, 23. 多机关门
|
|
# 24. 中心对称, 25. 部分对称, 26. 鱼骨
|
|
|
|
# 自动标注标签定义（暂时不用）：
|
|
# 0. 左右对称, 1. 上下对称, 2. 中心对称, 3. 斜向对称, 4. 伪对称, 5. 多房间, 6. 多走廊
|
|
# 32. 平面塔, 33. 转换塔, 34. 道具塔
|
|
|
|
# 标量值定义：
|
|
# 0. 整体密度，非空白图块/地图面积，空白图块还包括装饰图块
|
|
# 1. 墙体密度，墙壁/地图面积
|
|
# 2. 装饰密度，装饰数量/地图面积
|
|
# 3. 门密度，门数量/地图面积
|
|
# 4. 怪物密度，怪物数量/地图面积
|
|
# 5. 资源密度，资源数量/地图面积
|
|
# 6. 宝石密度，宝石数量/地图面积
|
|
# 7. 血瓶密度，血瓶数量/地图面积
|
|
# 8. 钥匙密度，钥匙数量/地图面积
|
|
# 9. 道具密度，道具数量/地图面积
|
|
# 10. 入口数量
|
|
# 11. 机关门数量
|
|
# 12. 咸鱼门数量（多层咸鱼门只算一个）
|
|
|
|
# 图块定义：
|
|
# 0. 空地, 1. 墙壁, 2. 门, 3. 钥匙, 4. 红宝石, 5. 蓝宝石, 6. 绿宝石, 7. 血瓶
|
|
# 8. 道具, 9. 怪物, 10. 入口, 14. 起始 token, 15. 终止 token
|
|
|
|
BATCH_SIZE = 8
|
|
LATENT_DIM = 48
|
|
KL_BETA = 0.1
|
|
SELF_GATE = 0.5
|
|
GATE_EPOCH = 5
|
|
VAL_BATCH_DIVIDER = 8
|
|
PROB_STEP = 0.05
|
|
NUM_CLASSES = 16
|
|
|
|
device = torch.device(
|
|
"cuda:1" if torch.cuda.is_available()
|
|
else "mps" if torch.mps.is_available()
|
|
else "cpu"
|
|
)
|
|
os.makedirs("result", exist_ok=True)
|
|
os.makedirs("result/vae", exist_ok=True)
|
|
os.makedirs("result/ginka_vae_img", exist_ok=True)
|
|
|
|
disable_tqdm = not sys.stdout.isatty()
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser(description="training codes")
|
|
parser.add_argument("--resume", type=bool, default=False)
|
|
parser.add_argument("--state_ginka", type=str, default="result/vae/ginka-100.pth")
|
|
parser.add_argument("--train", type=str, default="ginka-dataset.json")
|
|
parser.add_argument("--validate", type=str, default="ginka-eval.json")
|
|
parser.add_argument("--epochs", type=int, default=100)
|
|
parser.add_argument("--checkpoint", type=int, default=5)
|
|
parser.add_argument("--load_optim", type=bool, default=True)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
def train():
|
|
print(f"Using {device.type} to train model.")
|
|
|
|
args = parse_arguments()
|
|
|
|
vae = GinkaTransformerVAE(num_classes=NUM_CLASSES, latent_dim=LATENT_DIM).to(device)
|
|
|
|
dataset = GinkaRNNDataset(args.train, device)
|
|
dataset_val = GinkaRNNDataset(args.validate, device)
|
|
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
|
|
dataloader_val = DataLoader(dataset_val, batch_size=BATCH_SIZE // VAL_BATCH_DIVIDER, shuffle=True)
|
|
|
|
optimizer_ginka = optim.AdamW(vae.parameters(), lr=1e-4, weight_decay=1e-4)
|
|
# 自定义调度器允许在 self_prob 提高时重置调度器信息并提高学习率以适应学习
|
|
scheduler_ginka = VAEScheduler(
|
|
optimizer_ginka, factor=0.9, increase_factor=2, patience=10, max_lr=1e-4, min_lr=1e-6
|
|
)
|
|
|
|
criterion = VAELoss()
|
|
|
|
self_prob = 0
|
|
prob_epochs = 0
|
|
|
|
# 用于生成图片
|
|
tile_dict = dict()
|
|
for file in os.listdir('tiles'):
|
|
name = os.path.splitext(file)[0]
|
|
tile_dict[name] = cv2.imread(f"tiles/{file}", cv2.IMREAD_UNCHANGED)
|
|
|
|
if args.resume:
|
|
data_ginka = torch.load(args.state_ginka, map_location=device)
|
|
|
|
vae.load_state_dict(data_ginka["model_state"], strict=False)
|
|
|
|
if args.load_optim:
|
|
if data_ginka.get("optim_state") is not None:
|
|
optimizer_ginka.load_state_dict(data_ginka["optim_state"])
|
|
|
|
print("Train from loaded state.")
|
|
|
|
for epoch in tqdm(range(args.epochs), desc="VAE Training", disable=disable_tqdm):
|
|
loss_total = torch.Tensor([0]).to(device)
|
|
reco_loss_total = torch.Tensor([0]).to(device)
|
|
kl_loss_total = torch.Tensor([0]).to(device)
|
|
|
|
vae.teacher_forcing()
|
|
for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm):
|
|
target_map = batch["target_map"].to(device)
|
|
B, H, W = target_map.shape
|
|
input = target_map.view(B, H * W)
|
|
|
|
optimizer_ginka.zero_grad()
|
|
fake_logits, mu, logvar = vae(input, self_prob)
|
|
|
|
loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, input, mu, logvar, KL_BETA)
|
|
|
|
loss.backward()
|
|
torch.nn.utils.clip_grad_norm_(vae.parameters(), max_norm=1.0)
|
|
optimizer_ginka.step()
|
|
loss_total += loss.detach()
|
|
reco_loss_total += reco_loss.detach()
|
|
kl_loss_total += kl_loss.detach()
|
|
|
|
avg_loss = loss_total.item() / len(dataloader)
|
|
avg_reco_loss = reco_loss_total.item() / len(dataloader)
|
|
avg_kl_loss = kl_loss_total.item() / len(dataloader)
|
|
tqdm.write(
|
|
f"[Epoch {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " +
|
|
f"E: {epoch + 1} | Loss: {avg_loss:.6f} | Reco: {avg_reco_loss:.6f} | " +
|
|
f"KL: {avg_kl_loss:.6f} | Prob: {self_prob:.2f} | LR: {scheduler_ginka.get_last_lr()[0]:.6f}"
|
|
)
|
|
|
|
# 验证集
|
|
# with torch.no_grad():
|
|
# val_loss_total = torch.Tensor([0]).to(device)
|
|
# for batch in tqdm(dataloader_val, desc="Validating generator.", leave=False, disable=disable_tqdm):
|
|
# target_map = batch["target_map"].to(device)
|
|
|
|
# fake_logits, mu, logvar = vae(target_map, 1 - gt_prob)
|
|
|
|
# loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, target_map, mu, logvar, KL_BETA)
|
|
# val_loss_total += loss.detach()
|
|
|
|
# avg_loss_val = val_loss_total.item() / len(dataloader_val)
|
|
|
|
# 先使用训练集的损失值，因为过拟合比较严重，后续再想办法
|
|
if avg_loss < SELF_GATE:
|
|
prob_epochs += 1
|
|
else:
|
|
prob_epochs = 0
|
|
|
|
if prob_epochs >= GATE_EPOCH and self_prob < 1:
|
|
self_prob += PROB_STEP
|
|
prob_epochs = 0
|
|
if self_prob > 1:
|
|
self_prob = 1
|
|
|
|
self_prob = 1
|
|
|
|
scheduler_ginka.step(avg_loss, self_prob)
|
|
|
|
# 每若干轮输出一次图片，并保存检查点
|
|
if (epoch + 1) % 1 == 0:
|
|
# 保存检查点
|
|
torch.save({
|
|
"model_state": vae.state_dict(),
|
|
"optim_state": optimizer_ginka.state_dict(),
|
|
}, f"result/rnn/ginka-{epoch + 1}.pth")
|
|
|
|
val_loss_total = torch.Tensor([0]).to(device)
|
|
val_reco_loss_total = torch.Tensor([0]).to(device)
|
|
val_kl_loss_total = torch.Tensor([0]).to(device)
|
|
vae.eval()
|
|
with torch.no_grad():
|
|
idx = 0
|
|
gap = 5
|
|
color = (255, 255, 255) # 白色
|
|
vline = np.full((416, gap, 3), color, dtype=np.uint8) # 垂直分割线
|
|
# 地图重建展示
|
|
for batch in tqdm(dataloader_val, desc="Validating generator.", leave=False, disable=disable_tqdm):
|
|
target_map = batch["target_map"].to(device)
|
|
B, H, W = target_map.shape
|
|
input = target_map.view(B, H * W)
|
|
|
|
fake_logits, mu, logvar = vae(input, self_prob)
|
|
|
|
loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, input, mu, logvar, KL_BETA)
|
|
val_loss_total += loss.detach()
|
|
val_reco_loss_total += reco_loss.detach()
|
|
val_kl_loss_total += kl_loss.detach()
|
|
|
|
fake_map = torch.argmax(fake_logits, dim=2).view(B, H, W).cpu().numpy()
|
|
fake_img = matrix_to_image_cv(fake_map[0], tile_dict)
|
|
real_map = target_map.cpu().numpy()
|
|
real_img = matrix_to_image_cv(real_map[0], tile_dict)
|
|
img = np.block([[real_img], [vline], [fake_img]])
|
|
cv2.imwrite(f"result/ginka_vae_img/{idx}.png", img)
|
|
|
|
idx += 1
|
|
|
|
# 随机采样
|
|
for i in range(0, 8):
|
|
z = torch.randn(1, LATENT_DIM).to(device)
|
|
|
|
vae.autoregressive()
|
|
fake_logits = vae.decoder(z, torch.zeros(1, 169).to(device))
|
|
fake_map = fake_logits.view(-1, 13, 13).cpu().numpy()
|
|
fake_img = matrix_to_image_cv(fake_map[0], tile_dict)
|
|
|
|
cv2.imwrite(f"result/ginka_vae_img/{i}_rand.png", fake_img)
|
|
|
|
# 插值
|
|
val_length = len(dataset_val.data)
|
|
index1 = random.randint(0, val_length - 1)
|
|
index2 = random.randint(0, val_length - 1)
|
|
map1 = torch.LongTensor(dataset_val.data[index1]["map"]).to(device).view(1, 169)
|
|
map2 = torch.LongTensor(dataset_val.data[index2]["map"]).to(device).view(1, 169)
|
|
mu1, logvar1 = vae.encoder(map1)
|
|
mu2, logvar2 = vae.encoder(map2)
|
|
z1 = vae.reparameterize(mu1, logvar1)
|
|
z2 = vae.reparameterize(mu2, logvar2)
|
|
real_img1 = matrix_to_image_cv(map1[0].view(13, 13).cpu().numpy(), tile_dict)
|
|
real_img2 = matrix_to_image_cv(map2[0].view(13, 13).cpu().numpy(), tile_dict)
|
|
i = 0
|
|
for t in torch.linspace(0, 1, 8):
|
|
z = z1 * (1 - t / 8) + z2 * t / 8
|
|
fake_logits = vae.decoder(z, torch.zeros(1, 169).to(device))
|
|
fake_map = fake_logits.view(-1, 13, 13).cpu().numpy()
|
|
fake_img = matrix_to_image_cv(fake_map[0], tile_dict)
|
|
img = np.block([[real_img1], [vline], [fake_img], [vline], [real_img2]])
|
|
|
|
cv2.imwrite(f"result/ginka_vae_img/{i}_linspace.png", img)
|
|
i += 1
|
|
|
|
avg_loss_val = val_loss_total.item() / len(dataloader_val)
|
|
avg_reco_loss_val = val_reco_loss_total.item() / len(dataloader_val)
|
|
avg_kl_loss_val = val_kl_loss_total.item() / len(dataloader_val)
|
|
tqdm.write(
|
|
f"[Validate {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] E: {epoch + 1} | " +
|
|
f"Loss: {avg_loss_val:.6f} | Reco: {avg_reco_loss_val:.6f} | KL: {avg_kl_loss_val:.6f}"
|
|
)
|
|
|
|
print("Train ended.")
|
|
torch.save({
|
|
"model_state": vae.state_dict(),
|
|
}, f"result/ginka_transformer.pth")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
torch.set_num_threads(4)
|
|
train()
|