tpu_autoencoder.py

import os
import tensorflow as tf
from tensorflow.contrib import tpu
from tensorflow.contrib.cluster_resolver import TPUClusterResolver

IMAGE_FILES='gs://MY_OWN_GCS_BUCKET/celeba/files.txt'
SIZE=(128,128)
BATCH=128
LR=1e-4

def create_dataset():
    dataset = tf.data.TextLineDataset([IMAGE_FILES])
    dataset = (dataset.map(lambda x: tf.read_file(x))
            .map(lambda x: tf.image.decode_jpeg(x, channels=3))
            .map(lambda x: tf.to_float(x) / 255.)
            .map(lambda x: tf.image.resize_images(x, size=SIZE))
            .shuffle(2000)
            .batch(BATCH)
            )
    return dataset

def autoencoder(x):
    ch = 32
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch *=2
    x = tf.layers.conv2d(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch *=2
    x = tf.layers.conv2d(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch *=2
    x = tf.layers.conv2d(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch *=2
    x = tf.layers.conv2d(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu);
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch //=2
    x = tf.layers.conv2d_transpose(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch //=2
    x = tf.layers.conv2d_transpose(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch //=2
    x = tf.layers.conv2d_transpose(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, ch, 3, padding='same', activation=tf.nn.relu); ch //=2
    x = tf.layers.conv2d_transpose(x, ch, 2, strides=(2,2), padding='same', activation=tf.nn.relu)
    x = tf.layers.conv2d(x, 3, 3, padding='same', activation=tf.nn.sigmoid);
    return x

def autoencoder_loss(x, y):
    return tf.reduce_mean(tf.squared_difference(x,y))

dataset = create_dataset()
images_iterator = dataset.make_initializable_iterator()
images = images_iterator.get_next()
autoencoder_op = autoencoder(images)
loss_op = autoencoder_loss(images, autoencoder_op)

optimizer = tf.train.AdamOptimizer(learning_rate=LR)
train_op = optimizer.minimize(loss_op)

print(phase); phase += 1
tpu_grpc_url = TPUClusterResolver(tpu=[os.environ['TPU_NAME']]).get_master()

print(phase); phase += 1
sess = tf.Session(tpu_grpc_url)
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
sess.run(images_iterator.initializer)

for i in range(100):
    print(phase); phase += 1
    loss_out, _ = sess.run([loss_op, train_op])
    print('{}: {}'.format(i, loss_out))