From 948132797d919b0af8bad85a781c184b322b79d6 Mon Sep 17 00:00:00 2001 From: unanmed <1319491857@qq.com> Date: Fri, 6 Feb 2026 14:31:07 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E6=8D=A2=E5=9B=9E=20vae?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ginka/train_vae.py | 52 ++++++++++++++++++++++++++++++++-------- ginka/vae_rnn/encoder.py | 5 ++-- ginka/vae_rnn/loss.py | 8 +++++-- ginka/vae_rnn/vae.py | 9 +++---- 4 files changed, 56 insertions(+), 18 deletions(-) diff --git a/ginka/train_vae.py b/ginka/train_vae.py index ed3126e..9956e55 100644 --- a/ginka/train_vae.py +++ b/ginka/train_vae.py @@ -54,6 +54,7 @@ from shared.image import matrix_to_image_cv BATCH_SIZE = 128 LATENT_DIM = 48 +KL_BETA = 0.05 device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") os.makedirs("result", exist_ok=True) @@ -112,28 +113,51 @@ def train(): 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) for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm): target_map = batch["target_map"].to(device) optimizer_ginka.zero_grad() - fake_logits, z = vae(target_map, 1 - gt_prob) + fake_logits, mu, logvar = vae(target_map, 1 - gt_prob) - loss = criterion.vae_loss(fake_logits, target_map) + loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, target_map, 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} | " + - f"LR: {optimizer_ginka.param_groups[0]['lr']:.6f}" + f"E: {epoch + 1} | Loss: {avg_loss:.6f} | Reco: {avg_reco_loss:.6f} | " + + f"KL: {avg_kl_loss:.6f} | Prob: {gt_prob:.2f} | LR: {optimizer_ginka.param_groups[0]['lr']:.6f}" ) - scheduler_ginka.step(avg_loss) + # 验证集 + with torch.no_grad(): + 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 = criterion.vae_loss(fake_logits, target_map) + val_loss_total += loss.detach() + val_reco_loss_total += loss.detach() + val_kl_loss_total += loss.detach() + + idx += 1 + + if avg_loss_val < 0.5 and gt_prob > 0: + gt_prob -= 0.01 + + scheduler_ginka.step(avg_loss_val) # 每若干轮输出一次图片，并保存检查点 if (epoch + 1) % args.checkpoint == 0: @@ -144,6 +168,8 @@ def train(): }, 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) with torch.no_grad(): idx = 0 gap = 5 @@ -153,10 +179,12 @@ def train(): for batch in tqdm(dataloader_val, desc="Validating generator.", leave=False, disable=disable_tqdm): target_map = batch["target_map"].to(device) - fake_logits, z = vae(target_map, 1 - gt_prob) + fake_logits, mu, logvar = vae(target_map, 1 - gt_prob) - loss = criterion.vae_loss(fake_logits, target_map) + loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, target_map, 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=1).cpu().numpy() fake_img = matrix_to_image_cv(fake_map[0], tile_dict) @@ -183,8 +211,10 @@ def train(): index2 = random.randint(0, val_length - 1) map1 = torch.LongTensor(dataset_val.data[index1]["map"]).to(device).reshape(1, 13, 13) map2 = torch.LongTensor(dataset_val.data[index2]["map"]).to(device).reshape(1, 13, 13) - z1 = vae.encoder(map1) - z2 = vae.encoder(map2) + 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].cpu().numpy(), tile_dict) real_img2 = matrix_to_image_cv(map2[0].cpu().numpy(), tile_dict) i = 0 @@ -199,9 +229,11 @@ def train(): 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}" + f"Loss: {avg_loss_val:.6f} | Reco: {avg_reco_loss_val:.6f} | KL: {avg_kl_loss_val:.6f}" ) if avg_loss_val < 0.5 and gt_prob > 0: diff --git a/ginka/vae_rnn/encoder.py b/ginka/vae_rnn/encoder.py index ed1ac97..61332bf 100644 --- a/ginka/vae_rnn/encoder.py +++ b/ginka/vae_rnn/encoder.py @@ -121,8 +121,9 @@ class VAEEncoder(nn.Module): output[:, idx] = logits h = self.fusion(output) - vec = self.fc(h) - return vec + mu = self.fc_mu(h) + logvar = self.fc_logvar(h) + return mu, logvar if __name__ == "__main__": device = torch.device("cpu") diff --git a/ginka/vae_rnn/loss.py b/ginka/vae_rnn/loss.py index bebbc24..9595448 100644 --- a/ginka/vae_rnn/loss.py +++ b/ginka/vae_rnn/loss.py @@ -5,9 +5,13 @@ class VAELoss: def __init__(self): self.num_classes = 32 - def vae_loss(self, logits, target): + def vae_loss(self, logits, target, mu, logvar, beta=0.1): # target: [B, 13, 13] target = F.one_hot(target, num_classes=self.num_classes).float().permute(0, 3, 1, 2) recon_loss = F.cross_entropy(logits, target) - return recon_loss + kl_loss = -0.5 * torch.mean( + 1 + logvar - mu.pow(2) - logvar.exp() + ) + + return recon_loss + beta * kl_loss, recon_loss, kl_loss diff --git a/ginka/vae_rnn/vae.py b/ginka/vae_rnn/vae.py index ca2fbbd..d6550b9 100644 --- a/ginka/vae_rnn/vae.py +++ b/ginka/vae_rnn/vae.py @@ -18,9 +18,10 @@ class GinkaVAE(nn.Module): return mu + eps * std def forward(self, target_map: torch.Tensor, use_self_probility=0): - z = self.encoder(target_map) + mu, logvar = self.encoder(target_map) + z = self.reparameterize(mu, logvar) logits = self.decoder(z, target_map, use_self_probility) - return logits, z + return logits, mu, logvar if __name__ == "__main__": device = torch.device("cpu") @@ -34,13 +35,13 @@ if __name__ == "__main__": # 前向传播 start = time.perf_counter() - logits, z = model(input) + logits, mu, logvar = model(input) end = time.perf_counter() print_memory("前向传播后") print(f"推理耗时: {end - start}") - print(f"输出形状: logits= {logits.shape}, z={z.shape}") + print(f"输出形状: logits= {logits.shape}, mu={mu.shape}, logvar={logvar.shape}") print(f"Encoder parameters: {sum(p.numel() for p in model.encoder.parameters())}") print(f"Decoder parameters: {sum(p.numel() for p in model.decoder.parameters())}") print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")