├── .gitmodules
├── models
    ├── __init__.py
    ├── resnet_cifar.py
    └── util.py
├── get_data
    ├── __init__.py
    ├── __init__.pyc
    ├── get_data.pyc
    ├── download_and_convert_cifar_10.py
    └── util.py
├── slurm_tf_helper
    ├── __init__.py
    └── setup_clusters.py
├── train.sl
├── README.md
└── main.py


/.gitmodules:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/models/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/get_data/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/slurm_tf_helper/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/get_data/__init__.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/NERSC/cori-tf-distributed-examples/master/get_data/__init__.pyc


--------------------------------------------------------------------------------
/get_data/get_data.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/NERSC/cori-tf-distributed-examples/master/get_data/get_data.pyc


--------------------------------------------------------------------------------
/train.sl:
--------------------------------------------------------------------------------
 1 | #!/bin/bash -l
 2 | 
 3 | #SBATCH -p regular
 4 | #SBATCH -C haswell
 5 | #SBATCH -o batch_outputs/slurm_%N.%j.out
 6 | #SBATCH -e batch_outputs/slurm_%N.%j.out
 7 | #SBATCH --qos=premium
 8 | module load deeplearning
 9 | python get_data/download_and_convert_cifar_10.py
10 | srun python -u main.py $@
11 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # cori-tf-distributed-examples
 2 | Scripts/Benchmarks for Running Tensorflow Distributed on Cori
 3 | ### Running Code
 4 | #### Running In Batch
 5 | sbatch -N \<number of nodes\> -t \<time to run\> train.sl \<command line arguments to main.py\>
 6 | 
 7 | #### Running Interactively
 8 | salloc -N \<number of nodes\> -t \<time to run\> -C \<haswell or knl\>
 9 | 
10 | bash train.sl \<command line arguments to main.py\>
11 | 
12 | #### Command line arguments to main.py
13 |   
14 | usage: main.py [-h] [-m {async,sync}] [-b BATCHSIZE] [-p PATH_TO_H5]
15 | 
16 | optional arguments:
17 |   * -h, --help            show this help message and exit
18 |   
19 |   * -m {async,sync}, --mode {async,sync}
20 |                         which mode of distributed training to use: "sync" or
21 |                         "async" (default: sync)
22 |                         
23 |   * -b BATCHSIZE, --batchsize BATCHSIZE
24 |                         what batch size to use. That is, after each node gets
25 |                         a chunk of the data, how much data each node should
26 |                         process per iteration (default: 128)
27 |                         
28 |  * -p PATH_TO_H5, --path_to_h5 PATH_TO_H5
29 |                         path to hdf5 file for training (default: $SCRATCH/cifar10/cifar_10_caffe_hdf5/train.h5)
30 | 


--------------------------------------------------------------------------------
/models/resnet_cifar.py:
--------------------------------------------------------------------------------
 1 | # thanks to F. Chollet for this one:
 2 | # https://github.com/fchollet/deep-learning-models/blob/master/resnet50.py
 3 | 
 4 | 
 5 | import tensorflow as tf
 6 | from keras.layers import Conv2D, BatchNormalization, Activation, Input, MaxPooling2D, AveragePooling2D, Flatten, Dense
 7 | 
 8 | from util import conv_block, identity_block
 9 | import keras.backend as K
10 | 
11 | 
12 | def make_model(x_shape, batch_size=128, num_classes=10):
13 |     y = tf.placeholder(dtype=tf.int32,shape=(batch_size,))# Input(dtype=tf.int32, shape=y_shape)
14 |     K.set_learning_phase(1)
15 |     img_input = Input(batch_shape= tuple( [batch_size] + list(x_shape)))
16 |     bn_axis = 3
17 | 
18 | 
19 |     x = Conv2D(64, (3, 3), strides=(1, 1), padding='same',name='conv1')(img_input)
20 |     x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x)
21 |     x = Activation('relu')(x)
22 | 
23 |    
24 |     x = conv_block(x, 3, [16,16,16], stage=2, block='a', strides=(1, 1))
25 |     x = identity_block(x, 3, [16,16,16], stage=2, block='b')
26 |     x = identity_block(x, 3, [16,16,16], stage=2, block='c')
27 |     x = identity_block(x, 3, [16,16,16], stage=3, block='d')
28 | 
29 |     
30 |     x = conv_block(x, 3, [32,32,32], stage=3, block='a', strides=(2,2))
31 |     x = identity_block(x, 3, [32,32,32], stage=3, block='b')
32 |     x = identity_block(x, 3, [32,32,32], stage=3, block='c')
33 |     x = identity_block(x, 3, [32,32,32], stage=3, block='d')
34 | 
35 |     
36 |     
37 |     x = conv_block(x, 3, [64,64,64], stage=3, block='a', strides=(2,2))
38 |     x = identity_block(x, 3, [64,64,64], stage=3, block='b')
39 |     x = identity_block(x, 3, [64,64,64], stage=3, block='c')
40 |     x = identity_block(x, 3, [64,64,64], stage=3, block='d')
41 | 
42 |     
43 |     x = AveragePooling2D((8, 8), name='avg_pool')(x)
44 | 
45 | 
46 |     x = Flatten()(x)
47 |     x = Dense(num_classes, activation='softmax', name='fc1000')(x)
48 |     
49 |     loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=x,labels=y))
50 |     
51 |     return img_input, y, loss
52 | 
53 | 


--------------------------------------------------------------------------------
/slurm_tf_helper/setup_clusters.py:
--------------------------------------------------------------------------------
  1 | import tensorflow as tf
  2 | 
  3 | import sys
  4 | 
  5 | import os
  6 | 
  7 | def setup_slurm_cluster(num_ps=1):
  8 |     all_nodes = get_all_nodes()
  9 | 
 10 |     port = get_allowed_port()
 11 |     
 12 |     hostlist = [ ("%s:%i" % (node, port)) for node in all_nodes ] 
 13 |     ps_hosts, worker_hosts = get_parameter_server_and_worker_hosts(hostlist, num_ps=num_ps)
 14 | 
 15 | 
 16 |     proc_id, num_procs = get_slurm_proc_variables()
 17 |     
 18 |     num_tasks = num_procs - num_ps
 19 |     
 20 |     job_name = get_job_name(proc_id, num_ps)
 21 |     
 22 |     task_index = get_task_index(proc_id, job_name, num_ps)
 23 |     
 24 |     cluster_spec = make_cluster_spec(worker_hosts, ps_hosts)
 25 | 
 26 |     server = make_server(cluster_spec, job_name, task_index)
 27 |     
 28 |     return cluster_spec, server, task_index, num_tasks, job_name
 29 |     
 30 | 
 31 | def make_server(cluster_spec, job_name, task_index):
 32 |     server = tf.train.Server(cluster_spec,
 33 |                            job_name=job_name,
 34 |                            task_index=task_index)
 35 |     return server
 36 |     
 37 |     
 38 | def make_cluster_spec(worker_hosts, ps_hosts):
 39 |     cluster_spec = tf.train.ClusterSpec({
 40 |     "worker": worker_hosts,
 41 |     "ps": ps_hosts
 42 |     })
 43 |     return cluster_spec
 44 |     
 45 |     
 46 | def get_task_index(proc_id, job_name, num_ps):
 47 |     
 48 |     if job_name == "ps":
 49 |         task_index = proc_id
 50 |     elif job_name == "worker":
 51 |         #expects a task_index for workers that starts at 0
 52 |         task_index = proc_id - num_ps
 53 |     return task_index
 54 |     
 55 |     
 56 | def get_slurm_proc_variables():
 57 |     proc_id  = int( os.environ['SLURM_PROCID'] )
 58 |     num_procs     = int( os.environ['SLURM_NPROCS'] )
 59 |     return proc_id, num_procs
 60 |     
 61 | def get_job_name(proc_id, num_ps):
 62 |     if proc_id < num_ps:
 63 |         job_name = "ps"
 64 |     else:
 65 |         job_name = "worker"
 66 |     return job_name
 67 |     
 68 |     
 69 | def get_parameter_server_and_worker_hosts(hostlist, num_ps=1):
 70 |     """assumes num_ps nodes used for parameter server (one ps per node)
 71 |     and rest of nodes used for workers"""
 72 |     ps_hosts = hostlist[:num_ps]
 73 |     worker_hosts = hostlist[num_ps:]
 74 |     return ps_hosts, worker_hosts
 75 |     
 76 | def get_allowed_port():
 77 |     allowed_port = 22222
 78 |     return allowed_port
 79 | 
 80 | 
 81 | def get_all_nodes():
 82 |     return expand_nodelist( os.environ['SLURM_NODELIST'])
 83 |     
 84 | 
 85 | def expand_nodelist(node_string):
 86 |     try:
 87 |         pref, suff  = node_string.split('[')
 88 |     except:
 89 |         assert False, "Only one node allocated!"
 90 | 
 91 |     suff = suff. split(']')[0].split(',')
 92 |     nodes =[]
 93 |     for s in suff:
 94 |         if '-' not in s:
 95 |             nodes.append("%s%s" % (pref, s))
 96 |             continue
 97 |         beg,end = s.split('-')
 98 |         num_len=len(beg)
 99 |         for id in range(int(beg),int(end) + 1):
100 |             j= "%s%0" + str(num_len) + "d"
101 |             nodes.append(j % (pref, id))
102 |     return nodes
103 | 
104 | if __name__ == "__main__":
105 |     cluster, server, task_index, num_tasks, job_name = setup_slurm_cluster(num_ps=1)
106 | 


--------------------------------------------------------------------------------
/models/util.py:
--------------------------------------------------------------------------------
 1 | # thanks to F. Chollet for this one:
 2 | #https://github.com/fchollet/deep-learning-models/blob/master/resnet50.py
 3 | import tensorflow as tf
 4 | import os
 5 | from keras.layers import Conv2D, BatchNormalization, Activation, Input,Flatten, add
 6 | 
 7 | def identity_block(input_tensor, kernel_size, filters, stage, block):
 8 |     """The identity block is the block that has no conv layer at shortcut.
 9 |     # Arguments
10 |         input_tensor: input tensor
11 |         kernel_size: defualt 3, the kernel size of middle conv layer at main path
12 |         filters: list of integers, the filterss of 3 conv layer at main path
13 |         stage: integer, current stage label, used for generating layer names
14 |         block: 'a','b'..., current block label, used for generating layer names
15 |     # Returns
16 |         Output tensor for the block.
17 |     """
18 |     filters1, filters2, filters3 = filters
19 |     
20 |     bn_axis = 3
21 | 
22 |     conv_name_base = 'res' + str(stage) + block + '_branch'
23 |     bn_name_base = 'bn' + str(stage) + block + '_branch'
24 | 
25 |     x = Conv2D(filters1, (1, 1), name=conv_name_base + '2a')(input_tensor)
26 |     x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
27 |     x = Activation('relu')(x)
28 | 
29 |     x = Conv2D(filters2, kernel_size,
30 |                padding='same', name=conv_name_base + '2b')(x)
31 |     x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)
32 |     x = Activation('relu')(x)
33 | 
34 |     x = Conv2D(filters3, (1, 1), name=conv_name_base + '2c')(x)
35 |     x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x)
36 | 
37 |     x = add([x, input_tensor])
38 |     x = Activation('relu')(x)
39 |     return x
40 | 
41 | def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
42 |     """conv_block is the block that has a conv layer at shortcut
43 |     # Arguments
44 |         input_tensor: input tensor
45 |         kernel_size: defualt 3, the kernel size of middle conv layer at main path
46 |         filters: list of integers, the filterss of 3 conv layer at main path
47 |         stage: integer, current stage label, used for generating layer names
48 |         block: 'a','b'..., current block label, used for generating layer names
49 |     # Returns
50 |         Output tensor for the block.
51 |     Note that from stage 3, the first conv layer at main path is with strides=(2,2)
52 |     And the shortcut should have strides=(2,2) as well
53 |     """
54 |     filters1, filters2, filters3 = filters
55 |     bn_axis = 3
56 | 
57 |     conv_name_base = 'res' + str(stage) + block + '_branch'
58 |     bn_name_base = 'bn' + str(stage) + block + '_branch'
59 | 
60 |     x = Conv2D(filters1, (1, 1), strides=strides,
61 |                name=conv_name_base + '2a')(input_tensor)
62 |     x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
63 |     x = Activation('relu')(x)
64 | 
65 |     x = Conv2D(filters2, kernel_size, padding='same',
66 |                name=conv_name_base + '2b')(x)
67 |     x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)
68 |     x = Activation('relu')(x)
69 | 
70 |     x = Conv2D(filters3, (1, 1), name=conv_name_base + '2c')(x)
71 |     x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x)
72 | 
73 |     shortcut = Conv2D(filters3, (1, 1), strides=strides,
74 |                       name=conv_name_base + '1')(input_tensor)
75 |     shortcut = BatchNormalization(axis=bn_axis, name=bn_name_base + '1')(shortcut)
76 | 
77 |     x = add([x, shortcut])
78 |     x = Activation('relu')(x)
79 |     return x
80 | 
81 | 
82 | 
83 | 
84 | 
85 | 


--------------------------------------------------------------------------------
/get_data/download_and_convert_cifar_10.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from os.path import join
  3 | from subprocess import call
  4 | import numpy as np
  5 | from sklearn import model_selection
  6 | import h5py
  7 | 
  8 | def extract_cifar_data():
  9 |     user_scratch_path = os.environ["SCRATCH"]
 10 |     data_base_dir = join(user_scratch_path, "cori-tf-distributed-examples-data", "cifar10")
 11 | 
 12 | 
 13 |     if not os.path.exists(data_base_dir):
 14 |         os.makedirs(data_base_dir)
 15 | 
 16 |     tarpath = join(data_base_dir,"cifar-10-python.tar.gz")
 17 | 
 18 |     print("Downloading...")
 19 |     if not os.path.exists(tarpath):
 20 |         call(
 21 |             "wget --output-document %s http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"% tarpath,
 22 |             shell=True
 23 |         )
 24 |         print("Downloading done.\n")
 25 |     else:
 26 |         print("Dataset already downloaded. Did not download twice.\n")
 27 | 
 28 |     print("Extracting...")
 29 |     raw_cifar_directory = "cifar-10-batches-py"
 30 |     raw_data_path = join(data_base_dir, raw_cifar_directory)
 31 |     if not os.path.exists(raw_data_path):
 32 |         call(
 33 |             "cd %s; tar -zxvf %s" % (data_base_dir, tarpath),
 34 |             shell=True
 35 |         )
 36 |         print("Extracting successfully done to %s" % raw_data_path)
 37 |     else:
 38 |         print("Dataset already extracted. Did not extract twice.\n")
 39 | 
 40 |     cifar_directory = join(raw_data_path, "cifar_hdf5")
 41 |     train_filename = os.path.join(cifar_directory, 'train.h5')
 42 |     test_filename = os.path.join(cifar_directory, 'test.h5')
 43 |     if not os.path.exists(cifar_directory):
 44 | 
 45 |         def unpickle(file):
 46 |             try:
 47 |                 import cPickle as pickle
 48 |             except:
 49 |                 import pickle
 50 |             fo = open(file, 'rb')
 51 |             dict_ = pickle.load(fo)
 52 |             fo.close()
 53 |             return dict_
 54 | 
 55 |         def shuffle_data(data, labels):
 56 |             data, _, labels, _ = model_selection.train_test_split(
 57 |                 data, labels, test_size=0.0, random_state=42
 58 |             )
 59 |             return data, labels
 60 | 
 61 |         def load_data(train_batches):
 62 |             data = []
 63 |             labels = []
 64 |             for data_batch_i in train_batches:
 65 |                 d = unpickle(
 66 |                     os.path.join(raw_data_path, data_batch_i)
 67 |                 )
 68 |                 data.append(d['data'])
 69 |                 labels.append(np.array(d['labels']))
 70 |             # Merge training batches on their first dimension
 71 |             data = np.concatenate(data)
 72 |             labels = np.concatenate(labels)
 73 |             length = len(labels)
 74 | 
 75 |             data, labels = shuffle_data(data, labels)
 76 |             return data.reshape(length, 3, 32, 32), labels
 77 | 
 78 |         X, y = load_data(
 79 |             ["data_batch_{}".format(i) for i in range(1, 6)]
 80 |         )
 81 | 
 82 |         Xt, yt = load_data(["test_batch"])
 83 | 
 84 |         print("INFO: each dataset's element are of shape 3*32*32:")
 85 |         print('"print(X.shape)" --> "{}"\n'.format(X.shape))
 86 | 
 87 |         print("Data is fully loaded, now truly converting.")
 88 | 
 89 |         os.makedirs(cifar_directory)
 90 | 
 91 | 
 92 |         comp_kwargs = {'compression': 'gzip', 'compression_opts': 1}
 93 |         # Train
 94 |         with h5py.File(train_filename, 'w') as f:
 95 |             f.create_dataset('data', data=X, **comp_kwargs)
 96 |             f.create_dataset('label', data=y.astype(np.int_), **comp_kwargs)
 97 |         # Test
 98 |         with h5py.File(test_filename, 'w') as f:
 99 |             f.create_dataset('data', data=Xt, **comp_kwargs)
100 |             f.create_dataset('label', data=yt.astype(np.int_), **comp_kwargs)
101 | 
102 |         print('Conversion successfully done to "{}".\n'.format(cifar_directory))
103 |     else:
104 |         print("Conversion was already done. Did not convert twice.\n")
105 |     return train_filename, test_filename
106 | if __name__ == "__main__":
107 |     extract_cifar_data()


--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import os
  3 | 
  4 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  5 | parser.add_argument("-m", "--mode", default="sync", type=str,
  6 |                     help="which mode of distributed training to use: \"sync\" or \"async\"", 
  7 |                     choices=["async", "sync"])
  8 | 
  9 | parser.add_argument("-b","--batchsize", default=128, type=int,
 10 |                     help="what batch size to use. That is, after each node gets a chunk of the data, how much data each node should process per iteration")
 11 | parser.add_argument("-p", "--path_to_h5",help="path to hdf5 file for training",type=str,
 12 |                     default="%s/cori-tf-distributed-examples-data/cifar10/cifar-10-batches-py/cifar_hdf5/train.h5" % os.environ["SCRATCH"])
 13 | args = parser.parse_args()
 14 | 
 15 | 
 16 | import tensorflow as tf
 17 | import sys
 18 | from models.resnet_cifar import make_model
 19 | from slurm_tf_helper.setup_clusters import setup_slurm_cluster
 20 | import numpy as np
 21 | import time
 22 | import importlib
 23 | from get_data.util import get_data_generator
 24 | 
 25 | 
 26 | 
 27 | 
 28 | 
 29 | def main(_):
 30 |     
 31 |     num_classes = 10 #cuz cifar10
 32 |     cluster, server, task_index, num_tasks, job_name = setup_slurm_cluster(num_ps=1)
 33 |     
 34 |     if job_name == "ps":
 35 |         server.join()
 36 |     
 37 |     elif job_name == "worker":
 38 |         is_chief=(task_index == 0)
 39 |         
 40 |         #whatever you use to iterate over your data
 41 |         generator = get_data_generator(num_tasks=num_tasks, 
 42 |                                   task_id=task_index, 
 43 |                                   batch_size=args.batchsize, 
 44 |                                   path_to_h5=args.path_to_h5)
 45 |         
 46 |         
 47 |         
 48 |         # Replica Device Setter Assigns ops to the local worker by default and stores all variables in the parameter server (ps).
 49 |         device = tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % task_index,cluster=cluster)
 50 |         with tf.device(device):
 51 |             
 52 |             #set up model variables/ops (insert any tf or keras code here or whatever you use to define your model
 53 |             x,y, loss = make_model(generator.im_shape[1:], args.batchsize, num_classes)
 54 |             
 55 |             #global step that either gets updated after any node processes a batch (async) or when all nodes process a batch for a given iteration (sync)
 56 |             global_step = tf.contrib.framework.get_or_create_global_step()
 57 |             
 58 |             opt = tf.train.AdamOptimizer(1.0)
 59 |             
 60 |             if args.mode == "sync":
 61 |                 #if syncm we make a data structure that will aggregate the gradients form all tasks (one task per node in thsi case)
 62 |                 opt = tf.train.SyncReplicasOptimizer(opt, replicas_to_aggregate=num_tasks,
 63 |                                total_num_replicas=num_tasks)
 64 |             train_op = opt.minimize(loss, global_step=global_step)
 65 | 
 66 |         
 67 | 
 68 |         
 69 |         
 70 |         #a hook that will stop training at
 71 |         hooks=[tf.train.StopAtStepHook(last_step=1000000)]
 72 |         
 73 |         if args.mode == "sync":
 74 |             hooks.append(opt.make_session_run_hook(is_chief=is_chief))
 75 |             
 76 |         
 77 |         with tf.train.MonitoredTrainingSession(is_chief=is_chief,
 78 |                                                master=server.target,
 79 |                                                hooks=hooks) as mon_sess:
 80 |             
 81 |             
 82 |             while not mon_sess.should_stop():
 83 |                 
 84 |                 for ims,lbls in generator:
 85 |                     iteration_start = time.time()
 86 |                          
 87 |                     _, step, loss_ = mon_sess.run([train_op, global_step, loss],
 88 |                                                   feed_dict={x:ims, y: lbls})
 89 |                     print "Global Step is %i" % step
 90 |                     print "Loss for task id %i is: %8.4f " % (task_index, loss_)
 91 |                     print "Remember the iteration losses won't go down right away! Wait to see if the epoch losses go down, you eager beaver!"
 92 |                     print "Time for iteration for task %i, batch size %i is %5.2f" % (task_index, args.batchsize, time.time() - iteration_start)
 93 |                 
 94 | 
 95 | 
 96 | if __name__ == "__main__":
 97 |     tf.app.run(main=main)
 98 | 
 99 | 
100 | 
101 | 


--------------------------------------------------------------------------------
/get_data/util.py:
--------------------------------------------------------------------------------
  1 | import sys
  2 | import os
  3 | from os import listdir, system
  4 | from os.path import isfile, join, isdir
  5 | import numpy as np
  6 | import imp
  7 | import itertools
  8 | import os
  9 | import sys
 10 | import time
 11 | import inspect
 12 | import copy
 13 | import random
 14 | import collections
 15 | import h5py
 16 | import importlib
 17 | from download_and_convert_cifar_10 import extract_cifar_data
 18 | 
 19 | def get_data_generator(num_tasks, task_id,  path_to_h5, batch_size=128):
 20 |     ims, lbls = h5_data_splitter(path_to_h5, num_tasks, task_id)
 21 |     ims = np.transpose(ims,axes=(0,2,3,1))
 22 |     bf = BatchFetcher(ims, lbls)
 23 |     gen = ImGenerator(bf, batch_size=batch_size)
 24 |     return gen
 25 | 
 26 | def normalize(arr,min_=None, max_=None, axis=(0,2,3)):
 27 |     if min_ is None or max_ is None:
 28 |         min_ = arr.min(axis=(0,2,3), keepdims=True)
 29 | 
 30 |         max_ = arr.max(axis=(0,2,3), keepdims=True)
 31 | 
 32 |     midrange = (max_ + min_) / 2.
 33 | 
 34 |     range_ = (max_ - min_) / 2.
 35 |     #print range_
 36 |     arr -= midrange
 37 |     arr /= (range_)
 38 |     return arr
 39 | 
 40 | 
 41 | class ImGenerator(object):
 42 |     def __init__(self, batchfetcher, batch_size=128):
 43 |         self.data = batchfetcher
 44 |         self.batch_size = batch_size
 45 | 
 46 |     def __iter__(self):
 47 |         return self
 48 |    
 49 |     @property
 50 |     def num_ims(self):
 51 |         return self.data.num_examples
 52 |     
 53 |     @property
 54 |     def im_shape(self):
 55 |         return self.data.images.shape
 56 |     
 57 |     def next(self):
 58 | 
 59 |         ims,lbls = self.data.next_batch(batch_size=self.batch_size)
 60 |             
 61 |         return ims, lbls
 62 |     
 63 | def h5_data_splitter(path_to_h5file, num_tasks, task_id, image_key="data", label_key="label"):
 64 |     
 65 |     h5f = h5py.File(path_to_h5file)
 66 |     num_ims = h5f[image_key].shape[0]
 67 |     ims_per_task = int(num_ims / float(num_tasks) )
 68 |     task_slice = slice(task_id*ims_per_task, (task_id + 1)*ims_per_task)
 69 |     ims = h5f[image_key][task_slice]
 70 |     #ims = normalize(ims)
 71 |     lbls = h5f[label_key][task_slice]
 72 |     
 73 |     return ims, lbls
 74 |     
 75 | 
 76 | 
 77 | #Thanks to TensorFlow (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/learn/python/learn/datasets/mnist.py)
 78 | #for inspiration for this data structure
 79 | class BatchFetcher(object):
 80 | 
 81 |     def __init__(self, images, labels, num_examples=-1):
 82 |         """Create a Data Structure that can continually
 83 |        
 84 | 
 85 |  return minibatches of some data
 86 | 
 87 |         Keyword arguments:
 88 |         images -- an object that can be indexed (aka an array or a class with the __get_item__ fxn implemented)
 89 |                     * also the function shuffle and the attribute total_examples must be implemented
 90 |         labels -- same as above
 91 |         """
 92 | 
 93 | 
 94 |         self._images = images
 95 |         self._labels = labels
 96 |         self._num_examples = num_examples if num_examples != -1 else images.shape[0]
 97 |         self._epochs_completed = 0
 98 |         self._index_in_epoch = 0
 99 | 
100 |     @property
101 |     def images(self):
102 |         return self._images
103 | 
104 |     @property
105 |     def labels(self):
106 |         return self._labels
107 |     
108 |     @property
109 |     def num_examples(self):
110 |         return self._num_examples
111 | 
112 |     @property
113 |     def epochs_completed(self):
114 |         return self._epochs_completed
115 |     
116 |     def shuffle(self):
117 |         pass
118 | 
119 |     def next_batch(self, batch_size, shuffle=False):
120 |         """Return the next `batch_size` examples from this data set."""
121 |         start = self._index_in_epoch
122 |         # Shuffle for the first epoch
123 |         if self._epochs_completed == 0 and start == 0 and shuffle:
124 |             raise NotImplementedError
125 |             
126 |             #shuffle files
127 |             # no shuffle for now
128 | #             seed = np.random.randint(0,1000)
129 | #             self._images.shuffle(seed)
130 | #             self._labels.shuffle(seed)
131 | 
132 |         # Go to the next epoch
133 |         if start + batch_size > self._num_examples:
134 |             # Finished epoch
135 |             self._epochs_completed += 1
136 |             
137 |             # Get the rest examples in this epoch
138 |             rest_num_examples = self._num_examples - start
139 |             images_rest_part = self._images[start:self._num_examples]
140 |             labels_rest_part = self._labels[start:self._num_examples]
141 |             # Shuffle the data
142 |             if shuffle:
143 |                 raise NotImplementedError
144 |                 
145 | #                 seed = np.random.randint(0,1000)
146 | #                 self._images.shuffle(seed)
147 | #                 self._labels.shuffle(seed)
148 |             # Start next epoch
149 |             start = 0
150 |             self._index_in_epoch = batch_size - rest_num_examples
151 |             end = self._index_in_epoch
152 |             images_new_part = self._images[start:end]
153 |             labels_new_part = self._labels[start:end]
154 |             
155 | 
156 |             return np.concatenate((images_rest_part, images_new_part), axis=0),np.concatenate((labels_rest_part, labels_new_part), axis=0)
157 |         else:
158 |            
159 |             self._index_in_epoch += batch_size
160 |             end = self._index_in_epoch
161 |             images = self._images[start:end]
162 |             labels = self._labels[start:end]
163 |             #print images.shape
164 |             #print labels.shape
165 |             return images,labels
166 | 
167 | 
168 | 


--------------------------------------------------------------------------------