patch_count.py

from time import time
from glob import glob

import tensorflow as tf

from clustering.import_stage import build_id_import_stage


def main():

    with tf.Graph().as_default() as import_graph:
        all_files = glob('data/00238.tfrecord.gz')
        t_inputs = build_id_import_stage(all_files)

    coord = tf.train.Coordinator()
    with tf.Session(graph=import_graph) as sess:
        sess.run(tf.group(tf.local_variables_initializer(), tf.global_variables_initializer()))
        threads = tf.train.start_queue_runners(coord=coord, sess=sess)

        iteration = 0
        total_count = 0
        start_time = time()

        iid = 0
        pid = 0

        try:
            while not coord.should_stop():
                iteration += 1

                inputs = sess.run(t_inputs)
                iid = max(iid, inputs[0])
                pid = max(pid, inputs[1])

                batch_size = inputs.shape[0]

                total_count += batch_size

                if time() - start_time >= 10.:
                    start_time = time()
                    print('Current count: %i' % total_count)

        except tf.errors.OutOfRangeError:
            print('Training done after iteration %i' % (iteration-1))
        finally:
            coord.request_stop()
            coord.join(threads)

        print('Final count: %i' % total_count)
        print('IID %i. PID %i' % (iid, pid))

    pass


if __name__ == '__main__':
    main()