k/InfiniteMusic
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

made model biger

Browse Source
This commit is contained in:
k 2025-07-26 20:51:58 -04:00
parent a28901fdfd
commit 099400e1da
1 changed files with 13 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
model.py
View File

@ -8,9 +8,9 @@ class Model:
self.flattened_size = 128 * self.h * self.w self.flattened_size = 128 * self.h * self.w
# Encoder # Encoder
self.e1 = nn.Conv2d(input_channels, 32, kernel_size=3, stride=2, padding=1) self.e1 = nn.Conv2d(input_channels, 64, kernel_size=3, stride=2, padding=1)
self.e2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1) self.e2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
self.e3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1) self.e3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
# VAE Latent Space # VAE Latent Space
self.fc_mu = nn.Linear(self.flattened_size, latent_dim) self.fc_mu = nn.Linear(self.flattened_size, latent_dim)
@ -18,29 +18,29 @@ class Model:
# Decoder # Decoder
self.dl = nn.Linear(latent_dim, self.flattened_size) self.dl = nn.Linear(latent_dim, self.flattened_size)
self.d1 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1,output_padding=1) self.d1 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1,output_padding=1)
self.d2 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1,output_padding=1) self.d2 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1,output_padding=1)
self.d3 = nn.ConvTranspose2d(32, input_channels, kernel_size=3, stride=2, padding=1,output_padding=1) self.d3 = nn.ConvTranspose2d(64, input_channels, kernel_size=3, stride=2, padding=1,output_padding=1)
def __call__(self, x: Tensor) -> Tensor: def __call__(self, x: Tensor) -> Tensor:
mu, log_var = self.encode(x) mu, log_var = self.encode(x)
x = self.reparameterize(mu, log_var) x = self.reparameterize(mu, log_var)
return self.decode(x) return self.decode(x)
def __Lcall__(self, inp: Tensor, otp) -> (Tensor, Tensor): def __Lcall__(self, inp: Tensor, otp:Tensor, epoch) -> (Tensor, Tensor):
mu, log_var = self.encode(inp) mu, log_var = self.encode(inp)
z = self.reparameterize(mu, log_var) z = self.reparameterize(mu, log_var)
recon = self.decode(z) recon = self.decode(z)
# Normalized MSE (per-pixel) # Normalized MSE (per-pixel)
recon_loss = ((otp - recon).pow(2).mean()) recon_loss = (recon - otp).abs().mean()
# Stabilized KL # Stabilized KL
kl_div = -0.5 * (1 + log_var.clip(-10, 10) - mu.pow(2) - log_var.clip(-10, 10).exp()).mean() kl_div = -0.5 * (1 + log_var - mu.pow(2) - log_var.exp()).mean()
kl_div = kl_div.relu()
# Weighted loss # Weighted loss
total_loss = recon_loss + 0.5 * kl_div total_loss = recon_loss + min(0.1, 0.01 * epoch) * kl_div
return recon, total_loss return recon, total_loss
def encode(self, x: Tensor) -> (Tensor, Tensor): def encode(self, x: Tensor) -> (Tensor, Tensor):
@ -58,8 +58,8 @@ class Model:
def decode(self, x: Tensor) -> Tensor: def decode(self, x: Tensor) -> Tensor:
x = self.dl(x).relu() x = self.dl(x).relu()
x = x.reshape(shape=(-1, 128, self.h, self.w)) x = x.reshape(shape=(-1, 256, self.h, self.w))
x = self.d1(x).relu() x = self.d1(x).relu()
x = self.d2(x).relu() x = self.d2(x).relu()
x = self.d3(x).sigmoid() x = self.d3(x)
return x return x