File size: 2,448 Bytes
fd883a1 c77e6c8 |
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 |
from flax import linen as nn
import jax
import jax.numpy as jnp
from local_response_norm import LocalResponseNorm
LATENT_DIM = 500
EPSILON = 1e-8
class Generator(nn.Module):
@nn.compact
def __call__(self, latent, training=True):
x = nn.Dense(features=64)(latent)
# x = nn.BatchNorm(not training)(x)
x = nn.relu(x)
x = nn.Dense(features=2*2*1024)(x)
x = nn.BatchNorm(not training)(x)
x = nn.relu(x)
x = nn.Dropout(0.25, deterministic=not training)(x)
x = x.reshape((x.shape[0], 2, 2, -1))
x4 = nn.ConvTranspose(features=512, kernel_size=(3, 3), strides=(2, 2))(x)
x4 = LocalResponseNorm()(x4)
x4 = nn.relu(x4)
x4o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x4)
x4 = nn.ConvTranspose(features=512, kernel_size=(3, 3))(x4)
x4 = LocalResponseNorm()(x4)
x4 = nn.relu(x4)
x8 = nn.ConvTranspose(features=256, kernel_size=(3, 3), strides=(2, 2))(x4)
x8 = LocalResponseNorm()(x8)
x8 = nn.relu(x8)
x8o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x8)
x8 = nn.ConvTranspose(features=256, kernel_size=(3, 3))(x8)
x8 = LocalResponseNorm()(x8)
x8 = nn.relu(x8)
x16 = nn.ConvTranspose(features=128, kernel_size=(3, 3), strides=(2, 2))(x8)
x16 = LocalResponseNorm()(x16)
x16 = nn.relu(x16)
x16o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x16)
x16 = nn.ConvTranspose(features=128, kernel_size=(3, 3))(x16)
x16 = LocalResponseNorm()(x16)
x16 = nn.relu(x16)
x32 = nn.ConvTranspose(features=64, kernel_size=(3, 3), strides=(2, 2))(x16)
x32 = LocalResponseNorm()(x32)
x32 = nn.relu(x32)
x32o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x32)
x32 = nn.ConvTranspose(features=64, kernel_size=(3, 3))(x32)
x32 = LocalResponseNorm()(x32)
x32 = nn.relu(x32)
x64 = nn.ConvTranspose(features=32, kernel_size=(3, 3), strides=(2, 2))(x32)
x64 = LocalResponseNorm()(x64)
x64 = nn.relu(x64)
x64o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x64)
x64 = nn.ConvTranspose(features=32, kernel_size=(3, 3))(x64)
x64 = LocalResponseNorm()(x64)
x64 = nn.relu(x64)
x128 = nn.ConvTranspose(features=64, kernel_size=(3, 3), strides=(2, 2))(x64)
x128 = LocalResponseNorm()(x128)
x128 = nn.relu(x128)
x128o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x128)
return (nn.tanh(x128o), nn.tanh(x64o), nn.tanh(x32o), nn.tanh(x16o), nn.tanh(x8o), nn.tanh(x4o)) |