AncTe/ginka-generator
1
0
Fork 0
mirror of https://github.com/unanmed/ginka-generator.git synced 2026-05-20 09:11:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: 换回 vae

Browse Source
This commit is contained in:
unanmed 2026-02-06 14:31:07 +08:00
parent 9887abcd01
commit 948132797d
4 changed files with 56 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
ginka/train_vae.py
View File

@ -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:

5
ginka/vae_rnn/encoder.py
View File

@ -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")

8
ginka/vae_rnn/loss.py
View File

@ -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

9
ginka/vae_rnn/vae.py
View File

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