diff --git a/ginka/train_vae.py b/ginka/train_vae.py index 6a92426..fe66d4b 100644 --- a/ginka/train_vae.py +++ b/ginka/train_vae.py @@ -65,7 +65,12 @@ disable_tqdm = not sys.stdout.isatty() def gt_prob(epoch: int, max_epoch: int) -> float: progress = epoch / max_epoch - return 1 + if progress < 0.2: + return 1 + elif progress < 0.8: + return 1 - (progress - 0.2) / 0.6 + else: + return 0 def parse_arguments(): parser = argparse.ArgumentParser(description="training codes") @@ -115,30 +120,24 @@ 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) fake_logits, mu, logvar = vae(target_map, 1 - gt_prob(epoch, args.epochs)) - loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, target_map, mu, logvar, KL_BETA) + 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} | Reco Loss: {avg_reco_loss:.6f} | " + - f"KL Loss: {avg_kl_loss:.6f} | LR: {optimizer_ginka.param_groups[0]['lr']:.6f}" + f"E: {epoch + 1} | Loss: {avg_loss:.6f} | " + + f"LR: {optimizer_ginka.param_groups[0]['lr']:.6f}" ) scheduler_ginka.step() @@ -153,7 +152,6 @@ def train(): val_loss_total = torch.Tensor([0]).to(device) reco_loss_total = torch.Tensor([0]).to(device) - kl_loss_total = torch.Tensor([0]).to(device) with torch.no_grad(): idx = 0 gap = 5 @@ -163,12 +161,10 @@ 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, mu, logvar = vae(target_map, 1 - gt_prob(epoch, args.epochs)) + fake_logits, z = vae(target_map, 1 - gt_prob(epoch, args.epochs)) - loss, reco_loss, kl_loss = criterion.vae_loss(fake_logits, target_map, mu, logvar, KL_BETA) + loss = criterion.vae_loss(fake_logits, target_map, z, KL_BETA) val_loss_total += loss.detach() - reco_loss_total += reco_loss.detach() - 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) @@ -195,10 +191,8 @@ 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) - mu1, logvar1 = vae.encoder(map1) - mu2, logvar2 = vae.encoder(map2) - z1 = vae.reparameterize(mu1, logvar1) - z2 = vae.reparameterize(mu2, logvar2) + z1 = vae.encoder(map1) + z2 = vae.encoder(map2) 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 @@ -213,12 +207,9 @@ def train(): i += 1 avg_loss_val = val_loss_total.item() / len(dataloader_val) - avg_reco_loss = reco_loss_total.item() / len(dataloader_val) - avg_kl_loss = 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 Loss: {avg_reco_loss:.6f} | " + - f"KL Loss: {avg_kl_loss:.6f}" + f"Loss: {avg_loss_val:.6f}" ) print("Train ended.") diff --git a/ginka/vae_rnn/encoder.py b/ginka/vae_rnn/encoder.py index c42384e..ed1ac97 100644 --- a/ginka/vae_rnn/encoder.py +++ b/ginka/vae_rnn/encoder.py @@ -91,6 +91,9 @@ class VAEEncoder(nn.Module): self.embedding = EncoderEmbedding(tile_classes, width, height, 128, 256) self.rnn = EncoderGRU(256, self.rnn_hidden, self.logits_dim) self.fusion = EncoderFusion(256) + self.fc = nn.Sequential( + nn.Linear(512, latent_dim) + ) self.fc_mu = nn.Linear(512, latent_dim) self.fc_logvar = nn.Linear(512, latent_dim) @@ -118,9 +121,8 @@ class VAEEncoder(nn.Module): output[:, idx] = logits h = self.fusion(output) - mu = self.fc_mu(h) - logvar = self.fc_logvar(h) - return mu, logvar + vec = self.fc(h) + return vec if __name__ == "__main__": device = torch.device("cpu") diff --git a/ginka/vae_rnn/loss.py b/ginka/vae_rnn/loss.py index 9595448..bebbc24 100644 --- a/ginka/vae_rnn/loss.py +++ b/ginka/vae_rnn/loss.py @@ -5,13 +5,9 @@ class VAELoss: def __init__(self): self.num_classes = 32 - def vae_loss(self, logits, target, mu, logvar, beta=0.1): + def vae_loss(self, logits, target): # 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) - kl_loss = -0.5 * torch.mean( - 1 + logvar - mu.pow(2) - logvar.exp() - ) - - return recon_loss + beta * kl_loss, recon_loss, kl_loss + return recon_loss diff --git a/ginka/vae_rnn/vae.py b/ginka/vae_rnn/vae.py index d6550b9..ca2fbbd 100644 --- a/ginka/vae_rnn/vae.py +++ b/ginka/vae_rnn/vae.py @@ -18,10 +18,9 @@ class GinkaVAE(nn.Module): return mu + eps * std def forward(self, target_map: torch.Tensor, use_self_probility=0): - mu, logvar = self.encoder(target_map) - z = self.reparameterize(mu, logvar) + z = self.encoder(target_map) logits = self.decoder(z, target_map, use_self_probility) - return logits, mu, logvar + return logits, z if __name__ == "__main__": device = torch.device("cpu") @@ -35,13 +34,13 @@ if __name__ == "__main__": # 前向传播 start = time.perf_counter() - logits, mu, logvar = model(input) + logits, z = model(input) end = time.perf_counter() print_memory("前向传播后") print(f"推理耗时: {end - start}") - print(f"输出形状: logits= {logits.shape}, mu={mu.shape}, logvar={logvar.shape}") + print(f"输出形状: logits= {logits.shape}, z={z.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())}")