application_vnd.jupyter.stderr_output_1.png
from datalab import DataLabTrain

from vgg16 import vgg16

import tensorflow as tf

def train(n_iters):
    model, params = vgg16(fine_tune_last=True, n_classes=2)
    X = model['input']
    Z = model['out']
    Y = tf.placeholder(dtype=tf.float32, shape=[None, 2])
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=Z[:, 0, 0, :], labels=Y))
    train_step = tf.train.AdamOptimizer(0.0001).minimize(loss)
    saver = tf.train.Saver()
    with tf.Session() as sess:
        try:
            sess.run(tf.global_variables_initializer())
            for i in range(n_iters):
                dl = DataLabTrain('./datasets/train_set/')
                train_gen = dl.generator()
                dev_gen = DataLabTrain('./datasets/dev_set/').generator()
                for X_train, Y_train in train_gen:
                    sess.run(train_step, feed_dict={X: X_train, Y: Y_train})
                l = 0
                count = 0
                for X_test, Y_test in dev_gen:
                    count += 1
                    l += sess.run(loss, feed_dict={X: X_test, Y: Y_test})
                saver.save(sess, './model/vgg16-dog-vs-cat.ckpt')
        finally:
            sess.close()

train(n_iters=1)