Spaces:

lm-jim
/

mdl-mlops

Sleeping

lm-jim

Cambios finales para entrega

23e93fb 6 months ago

3.74 kB

import torch
import torch.nn as nn
import torch.nn.functional as F

import pytorch_lightning as pl

class ConvCVAE(pl.LightningModule):
    def __init__(self, latent_dim=20, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.latent_dim = latent_dim
        self.lr = lr

        # Red Encoder
        self.encoder_net = nn.Sequential(
            nn.Conv2d(11, 32, kernel_size=3, stride=2, padding=1), # 14x14
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 7x7
            nn.ReLU(),
            nn.Flatten(), # 64 * 7 * 7 = 3136
        )
        self.fc_mu = nn.Linear(3136, latent_dim)
        self.fc_logvar = nn.Linear(3136, latent_dim)

        # Red de proyeccion
        self.decoder_input = nn.Linear(latent_dim + 10, 3136)

        # Red Decoder
        self.decoder_net = nn.Sequential(
            nn.Unflatten(1, (64, 7, 7)), # 7x7
            nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1), # 14x14
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1), # 28x28
            nn.ReLU(),
            nn.Conv2d(32, 1, kernel_size=3, padding=1),
            nn.Sigmoid()
        )

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar) # Desviación estándar
        eps = torch.randn_like(std)   # Ruido aleatorio
        return mu + eps * std

    def forward(self, x, y):
        batch_size = x.size(0)

        # Codificamos con One-Hot el tag y lo expandimos a 28x28
        y_oh = F.one_hot(y, num_classes=10).float()
        y_map = y_oh.view(batch_size, 10, 1, 1).expand(batch_size, 10, 28, 28)

        # Aplicamos la red Encoder
        x_cond = torch.cat([x, y_map], dim=1) # 11 canales, 1 para la imagen y 10 para los tags
        x_enc = self.encoder_net(x_cond)
        mu, logvar = self.fc_mu(x_enc), self.fc_logvar(x_enc)

        # Parametrización
        z = self.reparameterize(mu, logvar)

        # Aplicamos la red Decoder
        z_cond = torch.cat([z, y_oh], dim=1) # Vector latente + tags

        # Red de proyección, reescalamos a 64x7x7
        d_in = self.decoder_input(z_cond)
        decoded_x = self.decoder_net(d_in)

        return decoded_x, mu, logvar

    def training_step(self, batch, batch_idx):
        x, y = batch
        recon_x, mu, logvar = self.forward(x, y)

        # Loss de reconstruccion (BCE)
        recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')

        # Loss de Divergencia KL
        kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

        # Promedio de ambos
        loss = (recon_loss + kld_loss) / x.size(0)

        self.log("train_loss", loss, prog_bar=True, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        recon_x, mu, logvar = self.forward(x, y)
        
        # Calculamos la misma pérdida
        recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')
        kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

        val_loss = (recon_loss + kld_loss) / x.size(0)
        
        self.log("val_loss", val_loss, prog_bar=True)
        return val_loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)
    
    def generate_number(self, number):
        label = torch.zeros(1, 10)
        label[:, number] = 1
        z = torch.randn(1, self.latent_dim)
        z_cond = torch.cat([z, label], dim=1)

        with torch.no_grad():
            d_in = self.decoder_input(z_cond)
            recon_x = self.decoder_net(d_in)
        return recon_x.squeeze().cpu().numpy()

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	import pytorch_lightning as pl

	class ConvCVAE(pl.LightningModule):
	def __init__(self, latent_dim=20, lr=1e-3):
	super().__init__()
	self.save_hyperparameters()
	self.latent_dim = latent_dim
	self.lr = lr

	# Red Encoder
	self.encoder_net = nn.Sequential(
	nn.Conv2d(11, 32, kernel_size=3, stride=2, padding=1), # 14x14
	nn.ReLU(),
	nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 7x7
	nn.ReLU(),
	nn.Flatten(), # 64 * 7 * 7 = 3136
	)
	self.fc_mu = nn.Linear(3136, latent_dim)
	self.fc_logvar = nn.Linear(3136, latent_dim)

	# Red de proyeccion
	self.decoder_input = nn.Linear(latent_dim + 10, 3136)

	# Red Decoder
	self.decoder_net = nn.Sequential(
	nn.Unflatten(1, (64, 7, 7)), # 7x7
	nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1), # 14x14
	nn.ReLU(),
	nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1), # 28x28
	nn.ReLU(),
	nn.Conv2d(32, 1, kernel_size=3, padding=1),
	nn.Sigmoid()
	)

	def reparameterize(self, mu, logvar):
	std = torch.exp(0.5 * logvar) # Desviación estándar
	eps = torch.randn_like(std) # Ruido aleatorio
	return mu + eps * std

	def forward(self, x, y):
	batch_size = x.size(0)

	# Codificamos con One-Hot el tag y lo expandimos a 28x28
	y_oh = F.one_hot(y, num_classes=10).float()
	y_map = y_oh.view(batch_size, 10, 1, 1).expand(batch_size, 10, 28, 28)

	# Aplicamos la red Encoder
	x_cond = torch.cat([x, y_map], dim=1) # 11 canales, 1 para la imagen y 10 para los tags
	x_enc = self.encoder_net(x_cond)
	mu, logvar = self.fc_mu(x_enc), self.fc_logvar(x_enc)

	# Parametrización
	z = self.reparameterize(mu, logvar)

	# Aplicamos la red Decoder
	z_cond = torch.cat([z, y_oh], dim=1) # Vector latente + tags

	# Red de proyección, reescalamos a 64x7x7
	d_in = self.decoder_input(z_cond)
	decoded_x = self.decoder_net(d_in)

	return decoded_x, mu, logvar

	def training_step(self, batch, batch_idx):
	x, y = batch
	recon_x, mu, logvar = self.forward(x, y)

	# Loss de reconstruccion (BCE)
	recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')

	# Loss de Divergencia KL
	kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

	# Promedio de ambos
	loss = (recon_loss + kld_loss) / x.size(0)

	self.log("train_loss", loss, prog_bar=True, on_epoch=True)
	return loss

	def validation_step(self, batch, batch_idx):
	x, y = batch
	recon_x, mu, logvar = self.forward(x, y)

	# Calculamos la misma pérdida
	recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')
	kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

	val_loss = (recon_loss + kld_loss) / x.size(0)

	self.log("val_loss", val_loss, prog_bar=True)
	return val_loss

	def configure_optimizers(self):
	return torch.optim.Adam(self.parameters(), lr=self.lr)

	def generate_number(self, number):
	label = torch.zeros(1, 10)
	label[:, number] = 1
	z = torch.randn(1, self.latent_dim)
	z_cond = torch.cat([z, label], dim=1)

	with torch.no_grad():
	d_in = self.decoder_input(z_cond)
	recon_x = self.decoder_net(d_in)
	return recon_x.squeeze().cpu().numpy()