tf_eager_minimize.py

import time

import tensorflow as tf

tf.enable_eager_execution()

# with tf.device('gpu:0'):
def model(xs):
  ys = tf.keras.layers.Conv2D(8, [2, 8], activation='relu')(xs)
  ys = tf.keras.layers.MaxPool2D([2, 2], strides=[2, 2])(ys)
  ys = tf.keras.layers.MaxPool2D([2, 2], strides=[2, 2])(ys)
  ys = tf.keras.layers.Conv2D(32, [2, 4], activation='relu')(ys)
  ys = tf.keras.layers.MaxPool2D([2, 2], strides=[2, 2])(ys)
  ys = tf.keras.layers.Conv2D(32, [2, 4], activation='relu')(ys)
  ys = tf.keras.layers.MaxPool2D([2, 2], strides=[1, 2])(ys)
  ys = tf.keras.layers.Flatten()(ys)
  ys = tf.keras.layers.Dense(2000, activation='relu')(ys)
  ys = tf.keras.layers.Dense(10, activation='softmax')(ys)
  return ys


def loss(xs, ys):
  return tf.keras.losses.categorical_crossentropy(model(xs), ys)

batch_size = 48
xs = tf.random_normal([batch_size, 43, 232, 1])
ys = tf.random_uniform([batch_size, 10])
# print(loss(xs, ys))
# print(model(xs).shape)

optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
print('Calling minimize...')
t0 = time.time()
for i in range(100):
  optimizer.minimize(lambda: loss(xs, ys))
print(time.time() - t0)