diff --git a/ginka/common/cond.py b/ginka/common/cond.py index 1b133fa..ddd7b44 100644 --- a/ginka/common/cond.py +++ b/ginka/common/cond.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from .common import DoubleFCModule class ConditionEncoder(nn.Module): - def __init__(self, tag_dim, val_dim, hidden_dim, out_dim): + def __init__(self, tag_dim=64, val_dim=16, hidden_dim=256, out_dim=256): super().__init__() self.tag_embed = DoubleFCModule(tag_dim, hidden_dim, hidden_dim) self.val_embed = DoubleFCModule(val_dim, hidden_dim, hidden_dim) diff --git a/ginka/dataset.py b/ginka/dataset.py index f4f8812..d811293 100644 --- a/ginka/dataset.py +++ b/ginka/dataset.py @@ -216,4 +216,28 @@ class GinkaWGANDataset(Dataset): return self.handle_stage4(target, tag_cond, val_cond) raise RuntimeError(f"Invalid train stage: {self.train_stage}") - \ No newline at end of file + +class GinkaRNNDataset(Dataset): + def __init__(self, data_path: str, device): + self.data = load_data(data_path) # 自定义数据加载函数 + self.device = device + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + item = self.data[idx] + + target = torch.LongTensor(item['map']) # [H, W] + H, W = target.shape + target = target.reshape(H * W) # [T] + tag_cond = torch.FloatTensor(item['tag']) + val_cond = torch.FloatTensor(item['val']) + val_cond[9] = val_cond[9] / H / W + val_cond[10] = val_cond[10] / H / W + + return { + "tag_cond": tag_cond, + "val_cond": val_cond, + "target_map": target + } \ No newline at end of file diff --git a/ginka/generator/loss.py b/ginka/generator/loss.py index 5755c92..6246700 100644 --- a/ginka/generator/loss.py +++ b/ginka/generator/loss.py @@ -401,3 +401,10 @@ class WGANGinkaLoss: ] return sum(losses) + +class RNNGinkaLoss: + def __init__(self): + pass + + def rnn_loss(self, fake, target): + return F.cross_entropy(fake, target) diff --git a/ginka/generator/rnn.py b/ginka/generator/rnn.py new file mode 100644 index 0000000..e40b55d --- /dev/null +++ b/ginka/generator/rnn.py @@ -0,0 +1,64 @@ +import time +import torch +import torch.nn as nn +import torch.nn.functional as F + +class GinkaRNN(nn.Module): + def __init__(self, tile_classes=32, cond_dim=256, input_dim=256, hidden_dim=512, num_layers=1): + super().__init__() + + # 输入部分 + self.embedding = nn.Embedding(tile_classes, input_dim) + self.input_fc = nn.Linear(input_dim, input_dim) + + self.gru = nn.GRU(input_dim + cond_dim, hidden_dim, num_layers, batch_first=True) + self.fc = nn.Linear(hidden_dim, tile_classes) + + def forward(self, x: torch.Tensor, cond: torch.Tensor): + """ + x: [B, T] + cond: [B, cond_dim] + """ + B, T = x.shape + tile_emb = self.input_fc(self.embedding(x)) # [B, T, input_dim] + cond_expand = cond.unsqueeze(1).expand(B, T, cond.shape[-1]) # [B, T, cond_dim] + + # 拼接 tile + cond + step_input = torch.cat([tile_emb, cond_expand], dim=-1) + + out, _ = self.gru(step_input) + logits = self.fc(out) + return logits + +def print_memory(tag=""): + print(f"{tag} | 当前显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB, 最大显存: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB") + +if __name__ == "__main__": + input = torch.rand(1, 32, 32, 32).cuda() + tag = torch.rand(1, 64).cuda() + val = torch.rand(1, 16).cuda() + + # 初始化模型 + model = GinkaRNN().cuda() + + print_memory("初始化后") + + # 前向传播 + start = time.perf_counter() + fake0 = model(input, 0, tag, val) + fake1 = model(F.softmax(fake0, dim=1), 1, tag, val) + fake2 = model(F.softmax(fake1, dim=1), 1, tag, val) + fake3 = model(F.softmax(fake2, dim=1), 1, tag, val) + end = time.perf_counter() + + print_memory("前向传播后") + + print(f"推理耗时: {end - start}") + print(f"输入形状: feat={input.shape}") + print(f"输出形状: output={fake3.shape}") + print(f"Random parameters: {sum(p.numel() for p in model.head.parameters())}") + print(f"Cond parameters: {sum(p.numel() for p in model.cond.parameters())}") + print(f"Input parameters: {sum(p.numel() for p in model.input.parameters())}") + print(f"UNet parameters: {sum(p.numel() for p in model.unet.parameters())}") + print(f"Output parameters: {sum(p.numel() for p in model.output.parameters())}") + print(f"Total parameters: {sum(p.numel() for p in model.parameters())}") diff --git a/ginka/train_rnn.py b/ginka/train_rnn.py new file mode 100644 index 0000000..28d2854 --- /dev/null +++ b/ginka/train_rnn.py @@ -0,0 +1,182 @@ +import argparse +import os +import sys +from datetime import datetime +import torch +import torch.optim as optim +import cv2 +from torch_geometric.loader import DataLoader +from tqdm import tqdm +from .common.cond import ConditionEncoder +from .generator.rnn import GinkaRNN +from .dataset import GinkaRNNDataset +from .generator.loss import RNNGinkaLoss +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-9. 黄蓝红门钥匙，机关门不使用钥匙开启 +# 10-12. 三种等级的红宝石 +# 13-15. 三种等级的蓝宝石 +# 16-18. 三种等级的绿宝石 +# 19-22. 四种等级的血瓶 +# 23-25. 三种等级的道具 +# 26-28. 三种等级的怪物 +# 29. 入口，不区分楼梯和箭头 + +BATCH_SIZE = 8 + +device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") +os.makedirs("result", exist_ok=True) +os.makedirs("result/rnn", 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/rnn/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 {'cuda' if torch.cuda.is_available() else 'cpu'} to train model.") + + args = parse_arguments() + + cond_inj = ConditionEncoder() + ginka_rnn = GinkaRNN() + + 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) + + scheduler_ginka = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer_ginka, T_0=10, T_mult=2) + optimizer_ginka = optim.Adam(list(ginka_rnn.parameters()) + list(cond_inj), lr=1e-3, betas=(0.0, 0.9)) + + criterion = RNNGinkaLoss() + + # 用于生成图片 + 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) + + ginka_rnn.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="RNN Training", disable=disable_tqdm): + loss_total_ginka = torch.Tensor([0]).to(device) + + iters = 0 + + for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm): + tag_cond = batch["tag_cond"].to(device) + val_cond = batch["val_cond"].to(device) + target_map = batch["target_map"].to(device) + + cond_vec = cond_inj(tag_cond, val_cond, 0) + fake = ginka_rnn(target_map, cond_vec) + + loss = criterion.rnn_loss(fake, target_map) + + loss.backward() + optimizer_ginka.step() + loss_total_ginka += loss.detach() + + iters += 1 + + if iters % 100 == 0: + avg_loss_ginka = loss_total_ginka.item() / iters + + tqdm.write( + f"[Iters {iters} {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " + + f"E: {epoch + 1} | Loss: {avg_loss_ginka:.6f} | " + + f"LR: {optimizer_ginka.param_groups[0]['lr']:.6f}" + ) + + avg_loss_ginka = loss_total_ginka.item() / iters + tqdm.write( + f"[Iters {iters} {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " + + f"E: {epoch + 1} | Loss: {avg_loss_ginka:.6f} | " + + f"LR: {optimizer_ginka.param_groups[0]['lr']:.6f}" + ) + + scheduler_ginka.step() + + # 每若干轮输出一次图片，并保存检查点 + if (epoch + 1) % args.checkpoint == 0: + # 保存检查点 + torch.save({ + "model_state": ginka_rnn.state_dict(), + "optim_state": optimizer_ginka.state_dict(), + }, f"result/wgan/ginka-{epoch + 1}.pth") + + val_loss_total = torch.Tensor([0]).to(device) + with torch.no_grad(): + idx = 0 + for batch in tqdm(dataloader_val, desc="Validating generator.", leave=False, disable=disable_tqdm): + tag_cond = batch["tag_cond"].to(device) + val_cond = batch["val_cond"].to(device) + target_map = batch["target_map"].to(device) + + cond_vec = cond_inj(tag_cond, val_cond, 0) + fake = ginka_rnn(target_map, cond_vec) + + val_loss_total += criterion.rnn_loss(fake, target_map).detach() + + fake_map = torch.argmax(fake, dim=1).cpu().numpy() + fake_img = matrix_to_image_cv(fake_map[0], tile_dict) + cv2.imwrite(f"result/ginka_rnn_img/{idx}.png", fake_img) + + idx += 1 + + print("Train ended.") + torch.save({ + "model_state": ginka_rnn.state_dict(), + }, f"result/ginka.pth") + + +if __name__ == "__main__": + torch.set_num_threads(4) + train()