<s> import numpy as np <EOL> class SupervisedDataset ( object ) : <EOL> def __init__ ( self , X , y ) : <EOL> self . X = X <EOL> self . n_samples = X . shape [ <NUM_LIT:0> ] <EOL> self . n_channels = X . shape [ <NUM_LIT:1> ] <EOL> self . height = X . shape [ <NUM_LIT:2> ] <EOL> self . width = X . shape [ <NUM_LIT:3> ] <EOL> self . y = y <EOL> def __iter__ ( self ) : <EOL> return self . iter <EOL> def iterator ( self , <EOL> mode = '<STR_LIT>' , <EOL> batch_size = None , <EOL> num_batches = None , <EOL> rng_seed = <NUM_LIT:0> ) : <EOL> if ( batch_size is not None ) : <EOL> if ( batch_size > self . n_samples ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if ( mode == '<STR_LIT>' ) : <EOL> if batch_size is None : <EOL> if num_batches is not None : <EOL> batch_size = int ( self . n_samples / num_batches ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> elif batch_size is not None : <EOL> if num_batches is not None : <EOL> max_num_batches = int ( np . ceil ( self . n_samples / batch_size ) ) <EOL> if num_batches > max_num_batches : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( self . n_samples , max_num_batches , <EOL> batch_size , num_batches ) ) <EOL> else : <EOL> num_batches = int ( self . n_samples / batch_size ) <EOL> self . iter = DatasetIteratorSequential ( self . X , self . y , batch_size , <EOL> num_batches , self . n_samples ) <EOL> elif ( mode == '<STR_LIT>' ) : <EOL> if batch_size is None : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> elif num_batches is None : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> self . iter = DatasetIteratorRandomUniform ( self . X , self . y , <EOL> batch_size , num_batches , <EOL> self . n_samples , rng_seed ) <EOL> elif ( mode == '<STR_LIT>' ) : <EOL> if batch_size is None : <EOL> if num_batches is not None : <EOL> batch_size = int ( np . floor ( np . float32 ( self . n_samples ) / <EOL> num_batches ) ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> elif batch_size is not None : <EOL> if num_batches is not None : <EOL> max_num_batches = int ( np . ceil ( self . n_samples / batch_size ) ) <EOL> if num_batches > max_num_batches : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( self . n_samples , max_num_batches , <EOL> batch_size , num_batches ) ) <EOL> else : <EOL> num_batches = int ( np . floor ( np . float32 ( self . n_samples ) / <EOL> batch_size ) ) <EOL> self . iter = DatasetIteratorRandomUniformNoRep ( self . X , self . y , <EOL> batch_size , <EOL> num_batches , <EOL> self . n_samples , <EOL> rng_seed ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> return self . iter <EOL> def get_num_samples ( self ) : <EOL> return self . n_samples <EOL> def get_batch ( self ) : <EOL> return self . iter . next ( ) <EOL> class BasicIterator ( object ) : <EOL> def __init__ ( self , X , y , <EOL> batch_size = None , <EOL> num_batches = None , <EOL> num_samples = None ) : <EOL> self . X = X <EOL> self . y = y <EOL> self . batch_size = batch_size <EOL> self . num_batches = num_batches <EOL> self . num_samples = num_samples <EOL> self . batch_count = <NUM_LIT:0> <EOL> self . sample_count = <NUM_LIT:0> <EOL> self . last = <NUM_LIT:0> <EOL> def reset ( self ) : <EOL> self . batch_count = <NUM_LIT:0> <EOL> self . sample_count = <NUM_LIT:0> <EOL> self . last = <NUM_LIT:0> <EOL> class DatasetIteratorSequential ( BasicIterator ) : <EOL> def __init__ ( self , X , y , <EOL> batch_size = None , <EOL> num_batches = None , <EOL> num_samples = None ) : <EOL> super ( DatasetIteratorSequential , self ) . __init__ ( X , y , <EOL> batch_size , <EOL> num_batches , <EOL> num_samples ) <EOL> def __iter__ ( self ) : <EOL> return self <EOL> def next ( self ) : <EOL> if ( self . batch_count >= self . num_batches or <EOL> self . sample_count >= self . num_samples ) : <EOL> raise StopIteration ( ) <EOL> elif ( self . sample_count + self . batch_size ) > self . num_samples : <EOL> self . last = slice ( self . sample_count , self . num_samples ) <EOL> self . sample_count = self . num_samples <EOL> return self . X [ self . last , : , : , : ] <EOL> else : <EOL> self . last = slice ( self . sample_count , <EOL> self . sample_count + self . batch_size ) <EOL> self . sample_count += self . batch_size <EOL> self . batch_count += <NUM_LIT:1> <EOL> return self . X [ self . last , : , : , : ] , self . y [ self . last ] <EOL> class DatasetIteratorRandomUniform ( BasicIterator ) : <EOL> def __init__ ( self , X , y , <EOL> batch_size = None , <EOL> num_batches = None , <EOL> num_samples = None , <EOL> rng_seed = <NUM_LIT:0> ) : <EOL> np . random . seed ( rng_seed ) <EOL> super ( DatasetIteratorRandomUniform , self ) . __init__ ( X , y , batch_size , <EOL> num_batches , <EOL> num_samples ) <EOL> def __iter__ ( self ) : <EOL> return self <EOL> def next ( self ) : <EOL> if self . batch_count >= self . num_batches : <EOL> raise StopIteration ( ) <EOL> else : <EOL> self . last = np . random . random_integers ( low = <NUM_LIT:0> , <EOL> high = self . num_samples - <NUM_LIT:1> , <EOL> size = ( self . batch_size , ) ) <EOL> self . batch_count += <NUM_LIT:1> <EOL> return self . X [ self . last , : , : , : ] , self . y [ self . last ] <EOL> class DatasetIteratorRandomUniformNoRep ( BasicIterator ) : <EOL> def __init__ ( self , X , y , <EOL> batch_size = None , <EOL> num_batches = None , <EOL> num_samples = None , <EOL> rng_seed = <NUM_LIT:0> ) : <EOL> np . random . seed ( rng_seed ) <EOL> self . remain_samples = np . arange ( num_samples ) <EOL> super ( DatasetIteratorRandomUniformNoRep , self ) . __init__ ( X , y , <EOL> batch_size , <EOL> num_batches , <EOL> num_samples ) <EOL> def __iter__ ( self ) : <EOL> return self <EOL> def next ( self ) : <EOL> if self . batch_count >= self . num_batches : <EOL> raise StopIteration ( ) <EOL> else : <EOL> idx = np . random . permutation ( <EOL> np . arange ( len ( self . remain_samples ) ) ) [ <NUM_LIT:0> : self . batch_size ] <EOL> self . last = self . remain_samples [ idx ] <EOL> self . remain_samples = np . setdiff1d ( self . remain_samples , self . last ) <EOL> self . batch_count += <NUM_LIT:1> <EOL> return self . X [ self . last , : , : , : ] , self . y [ self . last ] <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> X = np . zeros ( ( <NUM_LIT> , <NUM_LIT:3> , <NUM_LIT:5> , <NUM_LIT:5> ) ) <EOL> y = np . zeros ( ( <NUM_LIT> , ) ) <EOL> d = SupervisedDataset ( X , y ) <EOL> d . iterator ( mode = '<STR_LIT>' , batch_size = <NUM_LIT> ) <EOL> for i , [ b , b_l ] in enumerate ( d ) : <EOL> print ( '<STR_LIT>' . format ( i + <NUM_LIT:1> , b . shape ) ) <EOL> print ( '<STR_LIT>' . format ( b_l . shape ) ) </s>
