refactor: 换回 vae

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"E: {epoch + 1} | Loss: {avg_loss:.6f} | Reco: {avg_reco_loss:.6f} | " +
-            f"LR: {optimizer_ginka.param_groups[0]['lr']:.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)
+                mu1, logvar1 = vae.encoder(map1)
-                z2 = vae.encoder(map2)
+                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:

ginka/vae_rnn/encoder.py

@@ -121,8 +121,9 @@ class VAEEncoder(nn.Module):
             output[:, idx] = logits
 
         h = self.fusion(output)
-        vec = self.fc(h)
+        mu = self.fc_mu(h)
-        return vec
+        logvar = self.fc_logvar(h)
+        return mu, logvar
 
 if __name__ == "__main__":
     device = torch.device("cpu")

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

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())}")