File size: 4,658 Bytes
5e3f03f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | from typing import Any
import jax
import jax.numpy as jnp
from flax import linen as nn
Array = jax.Array
class ResBlock(nn.Module):
features: int
groups: int = 8
@nn.compact
def __call__(self, x: Array) -> Array:
residual = x
if x.shape[-1] != self.features:
residual = nn.Conv(self.features, kernel_size=(1, 1))(residual)
h = nn.GroupNorm(num_groups=min(self.groups, self.features))(x)
h = nn.swish(h)
h = nn.Conv(self.features, kernel_size=(3, 3))(h)
h = nn.GroupNorm(num_groups=min(self.groups, self.features))(h)
h = nn.swish(h)
h = nn.Conv(self.features, kernel_size=(3, 3))(h)
return h + residual
class Downsample(nn.Module):
features: int
@nn.compact
def __call__(self, x: Array) -> Array:
return nn.Conv(self.features, kernel_size=(3, 3), strides=(2, 2))(x)
class Upsample(nn.Module):
features: int
@nn.compact
def __call__(self, x: Array) -> Array:
b, h, w, _ = x.shape
x = jax.image.resize(x, (b, h * 2, w * 2, x.shape[-1]), method="nearest")
return nn.Conv(self.features, kernel_size=(3, 3))(x)
class Encoder(nn.Module):
latent_dim: int
base_channels: int = 64
channel_mults: tuple[int, ...] = (1, 2, 4)
@nn.compact
def __call__(self, x: Array) -> tuple[Array, Array]:
h = nn.Conv(self.base_channels, kernel_size=(3, 3))(x)
for i, mult in enumerate(self.channel_mults):
features = self.base_channels * mult
h = ResBlock(features)(h)
h = ResBlock(features)(h)
if i < len(self.channel_mults) - 1:
h = Downsample(features)(h)
h = Downsample(self.base_channels * self.channel_mults[-1])(h)
h = nn.GroupNorm(num_groups=8)(h)
h = nn.swish(h)
h = h.reshape(h.shape[0], -1)
mu = nn.Dense(self.latent_dim)(h)
log_var = nn.Dense(self.latent_dim)(h)
log_var = jnp.clip(log_var, -10.0, 10.0)
return mu, log_var
class Decoder(nn.Module):
out_channels: int
base_channels: int = 64
channel_mults: tuple[int, ...] = (1, 2, 4)
image_size: int = 32
@nn.compact
def __call__(self, z: Array) -> Array:
n_down = len(self.channel_mults)
start = self.image_size // (2**n_down)
features_top = self.base_channels * self.channel_mults[-1]
h = nn.Dense(start * start * features_top)(z)
h = h.reshape(z.shape[0], start, start, features_top)
h = Upsample(features_top)(h)
for i, mult in enumerate(reversed(self.channel_mults)):
features = self.base_channels * mult
h = ResBlock(features)(h)
h = ResBlock(features)(h)
if i < len(self.channel_mults) - 1:
next_features = self.base_channels * list(reversed(self.channel_mults))[i + 1]
h = Upsample(next_features)(h)
h = nn.GroupNorm(num_groups=8)(h)
h = nn.swish(h)
h = nn.Conv(self.out_channels, kernel_size=(3, 3))(h)
return nn.sigmoid(h)
class VAE(nn.Module):
latent_dim: int = 128
in_channels: int = 3
base_channels: int = 64
channel_mults: tuple[int, ...] = (1, 2, 4)
image_size: int = 32
obs_sigma: float = 0.1
def setup(self):
self.encoder = Encoder(
latent_dim=self.latent_dim,
base_channels=self.base_channels,
channel_mults=self.channel_mults,
)
self.decoder = Decoder(
out_channels=self.in_channels,
base_channels=self.base_channels,
channel_mults=self.channel_mults,
image_size=self.image_size,
)
def encode(self, x: Array) -> tuple[Array, Array]:
return self.encoder(x)
def decode(self, z: Array) -> Array:
return self.decoder(z)
def reparameterize(self, rng: Array, mu: Array, log_var: Array) -> Array:
eps = jax.random.normal(rng, mu.shape)
return mu + jnp.exp(0.5 * log_var) * eps
def __call__(self, x: Array, rng: Array) -> dict[str, Array]:
mu, log_var = self.encode(x)
z = self.reparameterize(rng, mu, log_var)
x_hat = self.decode(z)
return {"x_hat": x_hat, "mu": mu, "log_var": log_var, "z": z}
def sample_prior(self, rng: Array, n: int) -> Array:
z = jax.random.normal(rng, (n, self.latent_dim))
return self.decode(z)
def latent_log_prior(z: Array) -> Array:
"""Standard normal prior log density, summed over latent dimensions."""
d = z.shape[-1]
return -0.5 * (jnp.sum(z * z, axis=-1) + d * jnp.log(2 * jnp.pi))
|