Reflexiones y Jugando con Pixeles
Introducción
En mi post más reciente, me puse a jugar un poco con redes neuronales convolucionales y en particular con arquitecturas Codificador-Decodificador (Autocodificadores) para crear un simple motor de búsqueda de imágenes. Debo decir que, existen muchas razones por las cuales empecé estos proyectos personales, algunas por ejemplo:
- Familiarizarme más con
Pytorchya que ahora trabajaré en lado de infraestructura de este framework, por lo que necesito conocer sus usos y familiarizarme con laAPI. - Curiosidad, ya que en la práctica utilicé un modelo de calce de imágenes, y quise entender los fundamentos de fondo para este tipo de modelos.
- Aprender más sobre el estado del arte y los fundamentos que mueven la inteligencia artificial actualmente.
Jugando con Pixeles y Modelos Generativos
En esta sección simplente describiré algunos experimentos que hice para generar pixel-art (en particular sprites) de un estilo definido, a partir de un conjunto de datos de pixel art. Este mini-proyecto apareció por dos motivos:
- Un amigo me hizo volver al vicio de los video-juegos y como decimos en Conce (mi ciudad natal) me camellé 😅
- Tenía ganas de aprender un poco sobre lo reciente en redes neuronales
Inteligencia Artificial Generativa aplicada a imágenes
En un artículo anterior, hablé sobre cómo implementar un motor de búsqueda de imágenes. En esta sección, hablaré sobre un modelo de inteligencia artificial generativa (GenAI) que es una versión probabilística del modelo que explique previamente, y también intentaré resolver el problema de generar pixel-art a partir de un estilo definido en un conjunto de datos dado.
Autocodificador Variacional (VAE: Variational Auto-Encoder)
Asumimos que la variable $x$ que representa nuestros datos se genera a partir de una variable latente $z$ (representación codificada), la cual no es observable. Por lo tanto, el proceso generativo que cada dato sigue puede describirse como:
- Se hace un muestreo de la representación latente $z$ desde la distribución apriori $p(z)$
- Los datos originales, son muestrados de la distribución de probabilidad condicional $p(x|z)$
Con esta noción de modelo probabilistico, podemos definir una versión probabilística de los codificadores y decodificadores. El decodificador “probabilístico” está dado por $p(x|z)$ (obtenemos la reconstrucción de nuestros datos, dada su versión codificada), mientras que el “codificador probabilístico” está definido por $p(z|x)$, la cual describe la distribución de la variable codificada dada su versión decodificada.
Utilizando el Teorema de Bayes, podemos encontrar una relación entre estas distribuciones:
$$p(z|x) = \displaystyle \frac{p(x|z)p(z)}{p(x)} = \frac{p(x|z)p(z)}{\int p(x|u)p(u)du}$$
Ahora asumamos que $p(z)$ es una distribución Gaussiana y que $p(x|z)$ es también una distribución Gaussiana cuya media está definida por una función $f$ de la variable $z$ y cuya matriz de covarianza tiene la forma $cI$ donde $I$ es la matriz identidad y $c$ es una constante. Esta función $f$ pertenece a una familia de funciones $F$ que se deja sin especificar de momento y se escogerá más adelante. Hasta ahora, tenemos:
$$ \begin{array}{lll} p(z) \equiv \cal{N(0, I)} & & \\ p(x|z) \equiv \cal{N(f(z), cI)} & \quad f \in F & \quad c > 0 \end{array} $$
Consideremos que $f$ es fija y bien definida. Como mencionamos anteriormente, conocemos $p(z)$ y $p(x|z)$, por lo que podríamos utilizar el teorema de Bayes para calcular $p(z|x)$. Este es un problema de inferencia Bayesiana, que usualmente es intratable (integral en el denominador), y se requiere utilizar técnicas de aproximación.
En este caso, el problema puede ser visto como un problema de inferencia variacional, en el cual queremos aproximar $p(z|x)$ a una distribución Gaussiana $q_x(z)$, tal que su media y covarianza están definidas por dos funciones, $g$ y $h$, cuyo parámetro es $x$:
$$ \begin{array}{lll} q_x(z) \equiv \cal{N(g(x), h(x))} & \quad g \in G & \quad h \in H \end{array} $$
En simples términos, queremos minimizar la distancia entre estas dos distribuciones. Para ello podemos utilizar la divergencia de Kullback-Leibler entre la aproximación y la distribución $p(z|x)$ objetivo:
$$ \begin{align} (g^*, h^*) & = \underset{(g, h) \in G\times H}{\mathrm{argmin}} D_{KL}(q_x(z), p(z|x)) \\ & = \underset{(g, h) \in G\times H}{\mathrm{argmin}} \left(\mathop{\mathbb{E}_{z\sim q_x}}(\log q_x(z)) - \mathop{\mathbb{E}_{z\sim q_x}}\left(\log \displaystyle \frac{p(x|z)p(z)}{p(x)}\right) \right) \\ & = \underset{(g, h) \in G\times H}{\mathrm{argmin}} \left(\mathop{\mathbb{E}_{z\sim q_x}}(\log q_x(z)) - \mathop{\mathbb{E}_{z\sim q_x}}(\log p(z)) - \mathop{\mathbb{E}_{z\sim q_x}}(\log p(x|z)) + \mathop{\mathbb{E}_{z\sim q_x}}(\log p(x)) \right) \\ & = \underset{(g, h) \in G\times H}{\mathrm{argmax}} \left(\mathop{\mathbb{E}_{z\sim q_x}} (\log p(x|z)) - D_{KL}(q_x(z), p(z)) \right) \\ & = \underset{(g, h) \in G\times H}{\mathrm{argmax}} \left(\mathop{\mathbb{E}_{z\sim q_x}} \left(\displaystyle -\frac{||x - f(z)||}{2c}^2\right) - D_{KL}(q_x(z), p(z)) \right) \end{align} $$
Podemos identificar que existen dos términos, el error de reconstruccion entre $x$ y $f(z)$ y el término de regularización dado por la divergencia de Kullback-Leibler entre $q_x(z)$ y $p(z)$. Podemos también identificar la constante $c$ que balancea los dos términos. A mayor $c$ asumimos mayor varianza alrededor de $f(z)$.
Ahora si llevamos el modelo a redes neuronales, tendríamos una arquitectura como la mostrada en la figura 1.
Fig 1: Arquitectura codificador-decodificador variacional
Notar que el espacio latente no son puntos fijos, si no que cada componente mapea a una distribución de probabilidad. Notar que en medio de la red hay un proceso de muestreo. Este muestreo, debe hacerse tal que permita que el error se propague en la red neuronal (para actualizar los parámetros de la red).
Si simplemente muestreamos en medio de la red, va a ocurrir que agregamos aleatoriedad al proceso y el gradiente no va a poder fluir ya que será aleatorio en cada paso del algoritmo de retro-propagación. Un truco para evitar esto, es el truco de la re-parametrización. En este caso, dado que $z$ es una variable aleatoria que sigue una distribución Gaussiana, con media $g(x)$ y covarianza $H(x) = h(x) \cdot h^T(x)$, $z$ se puede expresar como:
$$z = h(x) \zeta + g(x) \quad \quad \zeta \sim \cal{N(0, I)}$$
Fig 2: Truco de la reparametrización
Finalmente, dada una imagen $x$ y su reconstrucción $\hat{x}$, entonces la función de costo que debe minimizarse para entrenar el codificador-decodificador variacional, puede escribirse como:
$$Loss = C ||x - \hat{x}||^2 + D_{KL}(\cal{N(\mu_x, \sigma_x)}, \cal{N(0, I)})$$
Generando Sprites (¡Pokemones!)
Por diversión, y para revivir años de vicio, quise intentar generar sprites de videojuegos. El conjunto de datos que utilicé es una lista de sprites de Pokémon.
Fig 4: Conjunto de datos utilizado.
Modelo VAE
En pytorch el codificador se vería como el siguiente código:
class Encoder(nn.Module):
def __init__(self, z_dim):
super().__init__()
self.z_dim = z_dim
self.model = nn.Sequential(
torch.nn.Conv2d(n_channels,
n_encoder_features,
kernel_size=(2, 2),
stride=(2, 2)),
nn.BatchNorm2d(n_encoder_features),
nn.ReLU(inplace=True),
torch.nn.Conv2d(n_encoder_features,
n_encoder_features * 2,
kernel_size=(2, 2),
stride=(2, 2)),
nn.BatchNorm2d(n_encoder_features * 2),
nn.ReLU(inplace=True),
torch.nn.Conv2d(n_encoder_features * 2,
n_encoder_features * 4,
kernel_size=3,
stride=1),
nn.BatchNorm2d(n_encoder_features * 4),
nn.ReLU(inplace=True),
torch.nn.Conv2d(n_encoder_features * 4,
n_encoder_features * 8,
kernel_size=3,
stride=1),
nn.BatchNorm2d(n_encoder_features * 8),
nn.ReLU(inplace=True),
)
# After all the convolutions we end up with a tensor of size 3x6
# Al finalizar las convoluciones con los parámetros definidos:
# (batch, n_encoder_features * 8, 3, 6)
self.flatten = torch.nn.Flatten()
self.dense = nn.Sequential(
torch.nn.Linear(8 * n_encoder_features * 3 * 6, 2 * z_dim),
)
def _reparametrize(self, mu, log_var):
zeta = torch.randn(*mu.shape, device=device)
return mu + torch.exp(log_var / 2) * zeta
def forward(self, x):
x = self.model(x)
x = self.flatten(x)
z = self.dense(x)
# De una capa obtenemos mu y log_var
mu, log_var = z[:,:self.z_dim], z[:,self.z_dim:]
z = self._reparametrize(mu, log_var)
return z, mu, log_var
Notar que al reparametrizar retorno el vector z muestreado, y los vectores de reparametrización $\mu$ y $\log \sigma^2$. La razón para trabajar con el logaritmo, es simplemente para la estabilidad del modelo. Si queremos convertir a la ecuación anterior, entonces tenemos que $\log \sigma^2 = 2 \log \sigma$, entonces torch.exp(log_var / 2) simplemente es $\sigma$.
El decodificador, simplemente toma este vector latente, y a partir de transformaciones (ej. convoluciones transpuestas), reconstruye la imagen original. Mi decodificador se ve algo así:
class Decoder(nn.Module):
def __init__(self, z_dim):
super().__init__()
self.model = nn.Sequential(
torch.nn.Linear(z_dim, 8 * n_encoder_features * 3 * 6),
torch.nn.Unflatten(1, (8 * n_encoder_features, 3, 6)),
torch.nn.ConvTranspose2d(8 * n_encoder_features,
n_decoder_features * 8,
kernel_size=1,
stride=1),
nn.BatchNorm2d(n_decoder_features * 8),
nn.ReLU(True),
torch.nn.ConvTranspose2d(8 * n_encoder_features,
n_decoder_features * 4,
kernel_size=2,
stride=2),
nn.BatchNorm2d(n_decoder_features * 4),
nn.ReLU(True),
torch.nn.ConvTranspose2d(n_decoder_features * 4,
n_decoder_features * 2,
kernel_size=(5, 9),
stride=(2, 1)),
nn.BatchNorm2d(n_decoder_features * 2),
nn.ReLU(True),
torch.nn.ConvTranspose2d(n_decoder_features * 2,
n_decoder_features,
kernel_size=(2, 2),
stride=(2, 2)),
nn.BatchNorm2d(n_decoder_features),
nn.ReLU(True),
torch.nn.ConvTranspose2d(n_decoder_features, n_channels, kernel_size=(1, 1), stride=1),
torch.nn.Tanh()
)
def forward(self, x):
return self.model(x)
Utilicé la función de activación $Tanh$, para que todos los elementos se encuentren entre -1 y 1 (mejor estabilidad y convergencia); Aunque podría haber utilizado una función sigmoide. Lo siguiente es definir la función de costo:
reconstruction_loss = torch.nn.MSELoss()
def vae_reconstruction_loss(y_true, y_pred, reconstruction_loss_factor):
return reconstruction_loss_factor * reconstruction_loss(y_true, y_pred)
def vae_kullback_leibler_loss(mu, log_var):
return -0.5*torch.sum(1 + log_var - mu**2 - torch.exp(log_var), axis=1)[0]
def vae_loss(y_true, y_pred, mu, log_var, reconstruction_loss_factor=1000):
recon_loss = vae_reconstruction_loss(y_true, y_pred, reconstruction_loss_factor)
kld_loss = vae_kullback_leibler_loss(mu, log_var)
return recon_loss + kld_loss
En el caso de la componente de reconstrucción, para hacer más claro el código la variable reconstruction_loss_factor representa la constante $C$ de la función de pérdida mostrada anteriormente. La función vae_kullback_leibler_loss calcula $D_{KL}(\cal{N(\mu_x, \sigma_x)}, \cal{N(0, I)})$:
$$ \begin{align} D_{KL}(\cal{N(\mu_x, \sigma_x)}, \cal{N(0, I)}) & = \mathop{\mathbb{E}} \left[\log \cal{N(\mu_x, \sigma_x)} - \log \cal{N(0, I)}\right] \\ & = \frac{1}{2} \left[\mu_x^2 + \sigma_x^2 - 1 - \log \sigma_x^2\right] \\\\ & = -\frac{1}{2} \left[1 + \log \sigma_x^2 - \mu_x^2 - \sigma_x^2 \right] \end{align} $$
Tuve que probar varios valores para el factor de error reconstrucción, al final los mejores resultados los obtuve con $C = 5000$. Para verificar el aprendizaje del modelo, fui monitoreando la función de costo en cada epoch. Entrené el modelo en 500 epochs.
Fig 4: Entrenamiento VAE en conjunto de datos.
Luego, tomé una muestra al azar, y inspeccioné las reconstrucciones que hace el modelo:
Fig 5: Ejemplo de datos en su versión original.
Fig 6: Ejemplo de datos reconstruidos por el VAE.
Finalmente, la parte divertida, generar Pokémones a partir de ruido:
Fig 7: Sprites de Pokémones generados a partir de muestreo en espacio latente y decodificador.
La verdad, para mi sigue siendo mágico que a partir de ruido se puedan generar nuevos sprites (pixel art). Cuando probé con el conjunto de datos de las caras de celebridades o el MNIST, también, en el espacio latente podía hacer operaciones lineales, y generar datos nuevos, similares a los datos de entrenamiento.
¿Es suficiente? Pixel-art es más complejo que imágenes convencionales
Podemos observar en los Pokémon generados anteriormente, que si bien tienen forma y coloreado similar a los datos vistos en el entrenamiento, estos siguen siendo borrosos y la representación considera un espacio continuo. Sin embargo, el pixel art y las imágenes en general, tienen un conjunto de colores limitados y en un espacio discreto.
Otro problema es que estamos intentando predecir todos los pixeles considerándolos como independientes, lo que es un supuesto demasiado simplista.
El pixel-art en general, utiliza una paleta de colores limitada, donde existen técnicas para coloreado, iluminación, aliasing que es diferente a la de una foto convencional (por ejemplo una fotografía “real”). Después de leer varios papers en el tema de pixel art, no encontré ninguna solución a mi problema. Sin embargo, investigando un poco más a fondo y expandiendo llegué a dos papers interesantes, que tratan los problemas con los que me topé:
- Pixel Recurrent Neural Networks
- Conditional Image Generation with PixelCNN Decoders
- PixelVAE: A Latent Variable Model for Natural Image
No quiero alargar el artículo más de la cuenta, así que no explicaré los modelos o los papers. En esencia, lo que intentan hacer estos modelos es, modelar el problema como un problema de predicción del siguiente pxiel (¿suena a algo parecido a lo que hacemos en NLP? 😊). Básicamente, queremos encontrar una distribución de probabilidad tal que:
$$p(x) = p(x_1, \ldots, x_n) = \displaystyle \prod_{i=1}^{n} p(x_i|x_1, \ldots, x_{i - 1})$$
Para lograr esto, se hace un modelamiento de imágenes autorregresivo, en este caso, el siguiente pixel, depende de los pixeles anteriores. Para ver detalles y un ejemplo de implementación simple, el siguiente tutorial es bastante completo:
La verdad, yo utilicé una pequeña variación del tutorial que acabo de mencionar.
Fig 8: Curva de aprendizaje de modelo PixelCNN con datos de sprites de Pokémon.
Intenté generar nuevos sprites con el modelo entrenado:
Fig 9: Pokémones generados con PixelCNN.
También intenté hacer autocompletado dada una parte de una imagen:
Fig 10: Autocompletado de Pokémones.
Para el muestreo, utilicé imágenes RGB, es decir 3 canales de color, y en este ejemplo se consideran los canales como independientes a la hora de muestrear. Esto en la práctica no es así, ya que el pixel-art utiliza una paleta de colores definida y limitada. La pregunta que me hago es ¿Existe alguna forma de considerar la paleta de colores en la entrada, para no tener que calcular 256 * n_canales de intensidades de pixeles posible?
La verdad sigo pensando e investigando cómo lograr esto, pero hasta ahora no he tenido buenos resultados 😔.
Observaciones sobre implementación y parámetros
Algunas observaciones que recalco y que me causaron curiosidad:
- La tasa de aprendizaje y el tamaño de los lotes (batches) influyen en la convergencia de la red
- Tuve explosión de gradiente cuando cualquiera de estos parámetros superaba ciertos umbrales.
- Óptimos locales dependiendo del parámetro de momentum.
- Normalización en la retro-propagación. Interesante, ya que sin normalizar también tuve problemas de explosión de gradiente y divergencia
- En la deconvolución repetir pixeles (Upsampling) o jugar con el kernel y distintos tamaño de saltos (Stride), también son muy dependientes del problema.
- Repetir las arquitecturas de los tutoriales, siempre resulta para ese caso específico y para los conjuntos de datos en las evaluaciones comparativas (Benchmark); ejemplo: CIFAR, MNIST, Celeb Faces, etc.
- Probé otras múltiples arquitecturas:
StyleGAN,StyleGAN2,Pix2Pix,CycleGAN. En el caso de las GAN tuve múltiples problemas como explosión de gradiente y desvanecimiento del mismo. Este tipo de redes son muy inestables al parecer. Por otro lado, las imágenes generadas no parecían pixel-art (manchas peores que las mostradas en este artículo 😂) - La cantidad de epochs es básicamente lo más importante, de ahí que el tener disponibilidad de GPU y poder de cómputo es un factor diferenciador.
Conclusiones
- VAE es un modelo generativo que intenta ajustar una distribución a cada punto en un espacio latente, lo que permite generar nuevos datos a partir de muestreo en dicho espacio.
- Generar pixel art es una tarea mucho más compleja que generar imágenes, debido a la naturaleza del pixel art: Paleta de colores limitada