import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import math
# A simple Block used in U-Net.
# down: x -> Conv -> Relu -> BN ->
# down: time -> dense -> relu -> add -> Conv -> relu -> BN -> transform
class Block(layers.Layer):
def __init__(self, in_ch, out_ch, time_emb_dim, up=False):
super(Block, self).__init__()
self.time_mlp = layers.Dense(out_ch)
if up:
self.conv1 = layers.Conv2D(out_ch, 3, padding='same')
self.transform = layers.Conv2DTranspose(out_ch, 4, strides=2, padding='same')
else:
self.conv1 = layers.Conv2D(out_ch, 3, padding='same')
self.transform = layers.Conv2D(out_ch, 4, strides=2, padding='same')
self.conv2 = layers.Conv2D(out_ch, 3, padding='same')
self.bnorm1 = layers.BatchNormalization()
self.bnorm2 = layers.BatchNormalization()
self.relu = layers.ReLU()
def call(self, x, t):
# First Conv
h = self.bnorm1(self.relu(self.conv1(x)))
# Time embedding
time_emb = self.relu(self.time_mlp(t))
# Extend last 2 dimensions
time_emb = tf.expand_dims(time_emb, axis=1)
time_emb = tf.expand_dims(time_emb, axis=1)
# Add time channel
h = h + time_emb
# Second Conv
h = self.bnorm2(self.relu(self.conv2(h)))
# Down or Upsample
return self.transform(h)
# Sinusoidal Position Embedding as used in Transformer
class SinusoidalPositionEmbeddings(layers.Layer):
def __init__(self, dim):
super(SinusoidalPositionEmbeddings, self).__init__()
self.dim = dim
def call(self, time):
device = time.device
time_float = tf.cast(time, dtype=tf.float32)
half_dim = self.dim // 2
embeddings = math.log(10000) / (half_dim - 1)
embeddings = tf.exp(tf.range(half_dim, dtype=tf.float32) * (-embeddings))
embeddings = time_float[:, None] * embeddings[None, :]
# embeddings = time_float * embeddings[None, :]
embeddings = tf.concat([tf.sin(embeddings), tf.cos(embeddings)], axis=-1)
return embeddings
class SimpleUnet(keras.Model):
def __init__(self):
super(SimpleUnet, self).__init__()
image_channels = 3
if DATASET == 'stanford_dogs':
down_channels = (IMG_SIZE, 128, 256, 512, 1024)
up_channels = (1024, 512, 256, 128, IMG_SIZE)
elif DATASET == 'cifar10':
down_channels = (IMG_SIZE, 128, 256)
up_channels = (256, 128, IMG_SIZE)
out_dim = 1
time_emb_dim = 32
# Time embedding
self.time_mlp = keras.Sequential([
SinusoidalPositionEmbeddings(time_emb_dim),
layers.Dense(time_emb_dim),
layers.ReLU()
])
# Initial projection
self.conv0 = layers.Conv2D(down_channels[0], 3, padding='same')
# Downsample
self.downs = [Block(down_channels[i], down_channels[i+1], time_emb_dim) for i in range(len(down_channels)-1)]
# Upsample
self.ups = [Block(up_channels[i], up_channels[i+1], time_emb_dim, up=True) for i in range(len(up_channels)-1)]
self.output_layer = layers.Conv2D(3, out_dim, padding='same')
def call(self, x, timestep):
# Embedd time
t = self.time_mlp(timestep)
# Initial conv
x = self.conv0(x)
# Unet
residual_inputs = []
for down in self.downs:
x = down(x, t)
residual_inputs.append(x)
for up in self.ups:
residual_x = residual_inputs.pop()
# Add residual x as additional channels
x = tf.concat([x, residual_x], axis=-1)
x = up(x, t)
return self.output_layer(x)
model = SimpleUnet()
