WorldModels-VAE-Network

VAE-Network.py

class Network(object):
	# Create model
	def __init__(self):
		self.image = tf.placeholder(tf.float32, [None, 96, 96, 3], name='image')
		self.resized_image = tf.image.resize_images(self.image, [64, 64])
		tf.summary.image('resized_image', self.resized_image, 20)

		self.z_mu, self.z_logvar = self.encoder(self.resized_image)
		self.z = self.sample_z(self.z_mu, self.z_logvar)
		self.reconstructions = self.decoder(self.z)
		tf.summary.image('reconstructions', self.reconstructions, 20)

		self.merged = tf.summary.merge_all()

		self.loss = self.compute_loss()

	def sample_z(self, mu, logvar):
		eps = tf.random_normal(shape=tf.shape(mu))
		return mu + tf.exp(logvar / 2) * eps

	def encoder(self, x):
		x = tf.layers.conv2d(x, filters=32, kernel_size=4, strides=2, padding='valid', activation=tf.nn.relu)
		x = tf.layers.conv2d(x, filters=64, kernel_size=4, strides=2, padding='valid', activation=tf.nn.relu)
		x = tf.layers.conv2d(x, filters=128, kernel_size=4, strides=2, padding='valid', activation=tf.nn.relu)
		x = tf.layers.conv2d(x, filters=256, kernel_size=4, strides=2, padding='valid', activation=tf.nn.relu)

		x = tf.layers.flatten(x)
		z_mu = tf.layers.dense(x, units=32, name='z_mu')
		z_logvar = tf.layers.dense(x, units=32, name='z_logvar')
		return z_mu, z_logvar

	def decoder(self, z):
		x = tf.layers.dense(z, 1024, activation=None)
		x = tf.reshape(x, [-1, 1, 1, 1024])
		x = tf.layers.conv2d_transpose(x, filters=128, kernel_size=5, strides=2, padding='valid', activation=tf.nn.relu)
		x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=5, strides=2, padding='valid', activation=tf.nn.relu)
		x = tf.layers.conv2d_transpose(x, filters=32, kernel_size=6, strides=2, padding='valid', activation=tf.nn.relu)
		x = tf.layers.conv2d_transpose(x, filters=3, kernel_size=6, strides=2, padding='valid', activation=tf.nn.sigmoid)
		return x