Update generator.py

generator.py

@@ -3,57 +3,47 @@ import jax
 import jax.numpy as jnp
 from local_response_norm import LocalResponseNorm
 
-LATENT_DIM = 500
 EPSILON = 1e-8
+MAX_DISC_FEATURES = 128
+MAX_GEN_FEATURES = 512
+LATENT_DIM = 512
+MAX_LAYERS = 7 # 256x256
+
+def get_gen_layers(layer):
+  resolution = int(4 * 2 ** layer)
+  features = min(int(32 * 2 ** (MAX_LAYERS - 1 - layer)), MAX_GEN_FEATURES)
+  layers = []
+  layers.append(lambda x: jax.image.resize(x, shape=(x.shape[0], resolution, resolution, x.shape[3]), method="linear"))
+  layers.append(lambda x: nn.ConvTranspose(features=features, kernel_size=(3, 3), name=f"ConvTranspose_{resolution}_{features}")(x))
+  layers.append(lambda x: nn.relu(x))
+  return layers
+
+def get_initial_gen_layers(num_layers):
+  layers = []
+  layers.append(lambda x: x.reshape(x.shape[0], 1, 1, -1))
+  return layers
+
+def get_final_gen_layers(num_layers):
+  resolution = int(4 * 2 ** (num_layers - 1))
+  layers = []
+  layers.append(lambda x: nn.ConvTranspose(features=3, kernel_size=(3, 3), name=f"ConvTranspose_{resolution}_3")(x))
+  return layers
 
 class Generator(nn.Module):
+  num_layers: int = None
+
+  def setup(self):
+    # 512 => 1x1x512 => 4x4x512 => 8x8x512 => 16x16x512 => 32x32x256 => 64x64x128 => 128x128x64 => 256x256x32 => 256x256x3
+    layers = []
+    layers.extend(get_initial_gen_layers(self.num_layers))
+    for layer in range(self.num_layers):
+      layers.extend(get_gen_layers(layer))
+    layers.extend(get_final_gen_layers(self.num_layers))
+    self.layers = layers
+
   @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))
+  def __call__(self, x):
+    result = x
+    for layer in self.layers:
+      result = layer(result)
+    return result