<s> import tensorflow as tf <EOL> import numpy as np <EOL> from collections import namedtuple <EOL> from tqdm import tqdm <EOL> from six . moves import zip , range <EOL> import multiprocessing <EOL> from . utils . concurrency import ensure_proc_terminate , OrderedResultGatherProc , DIE <EOL> from . tfutils import * <EOL> from . utils import logger <EOL> from . tfutils . modelutils import describe_model <EOL> from . dataflow import DataFlow , BatchData <EOL> from . dataflow . dftools import dataflow_to_process_queue <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> PredictResult = namedtuple ( '<STR_LIT>' , [ '<STR_LIT:input>' , '<STR_LIT>' ] ) <EOL> class PredictConfig ( object ) : <EOL> def __init__ ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> def assert_type ( v , tp ) : <EOL> assert isinstance ( v , tp ) , v . __class__ <EOL> self . session_config = kwargs . pop ( '<STR_LIT>' , get_default_sess_config ( ) ) <EOL> assert_type ( self . session_config , tf . ConfigProto ) <EOL> self . session_init = kwargs . pop ( '<STR_LIT>' ) <EOL> self . model = kwargs . pop ( '<STR_LIT>' ) <EOL> self . input_data_mapping = kwargs . pop ( '<STR_LIT>' , None ) <EOL> self . output_var_names = kwargs . pop ( '<STR_LIT>' ) <EOL> self . nr_gpu = kwargs . pop ( '<STR_LIT>' , <NUM_LIT:1> ) <EOL> assert len ( kwargs ) == <NUM_LIT:0> , '<STR_LIT>' . format ( str ( kwargs . keys ( ) ) ) <EOL> def get_predict_func ( config ) : <EOL> """<STR_LIT>""" <EOL> output_var_names = config . output_var_names <EOL> input_vars = config . model . get_input_vars ( ) <EOL> cost_var = config . model . get_cost ( input_vars , is_training = False ) <EOL> if config . input_data_mapping is None : <EOL> input_map = input_vars <EOL> else : <EOL> input_map = [ input_vars [ k ] for k in config . input_data_mapping ] <EOL> output_vars = [ tf . get_default_graph ( ) . get_tensor_by_name ( get_op_var_name ( n ) [ <NUM_LIT:1> ] ) <EOL> for n in output_var_names ] <EOL> sess = tf . Session ( config = config . session_config ) <EOL> config . session_init . init ( sess ) <EOL> def run_input ( dp ) : <EOL> assert len ( input_map ) == len ( dp ) , "<STR_LIT>" . format ( <EOL> len ( input_map ) , len ( dp ) ) <EOL> feed = dict ( zip ( input_map , dp ) ) <EOL> return sess . run ( output_vars , feed_dict = feed ) <EOL> return run_input <EOL> class PredictWorker ( multiprocessing . Process ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , idx , gpuid , inqueue , outqueue , config ) : <EOL> """<STR_LIT>""" <EOL> super ( PredictWorker , self ) . __init__ ( ) <EOL> self . idx = idx <EOL> self . gpuid = gpuid <EOL> self . inqueue = inqueue <EOL> self . outqueue = outqueue <EOL> self . config = config <EOL> def run ( self ) : <EOL> os . environ [ '<STR_LIT>' ] = self . gpuid <EOL> G = tf . Graph ( ) <EOL> with G . as_default ( ) , tf . device ( '<STR_LIT>' ) : <EOL> self . func = get_predict_func ( self . config ) <EOL> if self . idx == <NUM_LIT:0> : <EOL> describe_model ( ) <EOL> while True : <EOL> tid , dp = self . inqueue . get ( ) <EOL> if tid == DIE : <EOL> self . outqueue . put ( ( DIE , None ) ) <EOL> return <EOL> else : <EOL> res = PredictResult ( dp , self . func ( dp ) ) <EOL> self . outqueue . put ( ( tid , res ) ) <EOL> class DatasetPredictor ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , config , dataset ) : <EOL> """<STR_LIT>""" <EOL> assert isinstance ( dataset , DataFlow ) <EOL> self . ds = dataset <EOL> self . nr_gpu = config . nr_gpu <EOL> if self . nr_gpu > <NUM_LIT:1> : <EOL> self . inqueue , self . inqueue_proc = dataflow_to_process_queue ( self . ds , <NUM_LIT:10> , self . nr_gpu ) <EOL> self . outqueue = multiprocessing . Queue ( ) <EOL> try : <EOL> gpus = os . environ [ '<STR_LIT>' ] . split ( '<STR_LIT:U+002C>' ) <EOL> except KeyError : <EOL> gpus = list ( range ( self . nr_gpu ) ) <EOL> self . workers = [ PredictWorker ( i , gpus [ i ] , self . inqueue , self . outqueue , config ) <EOL> for i in range ( self . nr_gpu ) ] <EOL> self . result_queue = OrderedResultGatherProc ( self . outqueue ) <EOL> self . inqueue_proc . start ( ) <EOL> for p in self . workers : p . start ( ) <EOL> self . result_queue . start ( ) <EOL> ensure_proc_terminate ( self . workers ) <EOL> ensure_proc_terminate ( [ self . result_queue , self . inqueue_proc ] ) <EOL> else : <EOL> self . func = get_predict_func ( config ) <EOL> def get_result ( self ) : <EOL> """<STR_LIT>""" <EOL> with tqdm ( total = self . ds . size ( ) ) as pbar : <EOL> if self . nr_gpu == <NUM_LIT:1> : <EOL> for dp in self . ds . get_data ( ) : <EOL> yield PredictResult ( dp , self . func ( dp ) ) <EOL> pbar . update ( ) <EOL> else : <EOL> die_cnt = <NUM_LIT:0> <EOL> while True : <EOL> res = self . result_queue . get ( ) <EOL> if res [ <NUM_LIT:0> ] != DIE : <EOL> yield res [ <NUM_LIT:1> ] <EOL> else : <EOL> die_cnt += <NUM_LIT:1> <EOL> if die_cnt == self . nr_gpu : <EOL> break <EOL> pbar . update ( ) <EOL> self . inqueue_proc . join ( ) <EOL> self . inqueue_proc . terminate ( ) <EOL> for p in self . workers : <EOL> p . join ( ) ; p . terminate ( ) <EOL> def get_all_result ( self ) : <EOL> """<STR_LIT>""" <EOL> return list ( self . get_result ( ) ) </s>
