<s> import tensorflow as tf <EOL> import threading <EOL> import copy <EOL> import re <EOL> import functools <EOL> from six . moves import zip <EOL> from . base import Trainer <EOL> from . . dataflow . common import RepeatedData <EOL> from . . utils import * <EOL> from . . utils . concurrency import LoopThread <EOL> from . . tfutils . summary import summary_moving_average <EOL> from . . tfutils import * <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> class SimpleTrainer ( Trainer ) : <EOL> def run_step ( self ) : <EOL> data = next ( self . data_producer ) <EOL> feed = dict ( zip ( self . input_vars , data ) ) <EOL> self . sess . run ( [ self . train_op ] , feed_dict = feed ) <EOL> def train ( self ) : <EOL> model = self . model <EOL> input_vars = model . get_input_vars ( ) <EOL> self . input_vars = input_vars <EOL> cost_var = model . get_cost ( input_vars , is_training = True ) <EOL> avg_maintain_op = summary_moving_average ( cost_var ) <EOL> grads = self . config . optimizer . compute_gradients ( cost_var ) <EOL> grads = self . process_grads ( grads ) <EOL> self . train_op = tf . group ( <EOL> self . config . optimizer . apply_gradients ( grads , get_global_step_var ( ) ) , <EOL> avg_maintain_op ) <EOL> self . init_session_and_coord ( ) <EOL> self . data_producer = RepeatedData ( self . config . dataset , - <NUM_LIT:1> ) . get_data ( ) <EOL> self . main_loop ( ) <EOL> def _trigger_epoch ( self ) : <EOL> if self . summary_op is not None : <EOL> data = next ( self . data_producer ) <EOL> feed = dict ( zip ( self . input_vars , data ) ) <EOL> summary_str = self . summary_op . eval ( feed_dict = feed ) <EOL> self . _process_summary ( summary_str ) <EOL> class EnqueueThread ( threading . Thread ) : <EOL> def __init__ ( self , trainer , queue , enqueue_op , raw_input_var ) : <EOL> super ( EnqueueThread , self ) . __init__ ( ) <EOL> self . sess = trainer . sess <EOL> self . coord = trainer . coord <EOL> self . dataflow = RepeatedData ( trainer . config . dataset , - <NUM_LIT:1> ) <EOL> self . input_vars = raw_input_var <EOL> self . op = enqueue_op <EOL> self . queue = queue <EOL> self . close_op = self . queue . close ( cancel_pending_enqueues = True ) <EOL> self . daemon = True <EOL> def run ( self ) : <EOL> try : <EOL> while True : <EOL> for dp in self . dataflow . get_data ( ) : <EOL> if self . coord . should_stop ( ) : <EOL> return <EOL> feed = dict ( zip ( self . input_vars , dp ) ) <EOL> self . op . run ( feed_dict = feed , session = self . sess ) <EOL> except tf . errors . CancelledError as e : <EOL> pass <EOL> except Exception : <EOL> logger . exception ( "<STR_LIT>" ) <EOL> self . sess . run ( self . close_op ) <EOL> self . coord . request_stop ( ) <EOL> finally : <EOL> logger . info ( "<STR_LIT>" ) <EOL> class QueueInputTrainer ( Trainer ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , config , input_queue = None , async = False ) : <EOL> """<STR_LIT>""" <EOL> super ( QueueInputTrainer , self ) . __init__ ( config ) <EOL> self . input_vars = self . model . get_input_vars ( ) <EOL> if input_queue is None : <EOL> self . input_queue = tf . FIFOQueue ( <EOL> <NUM_LIT:100> , [ x . dtype for x in self . input_vars ] , name = '<STR_LIT>' ) <EOL> else : <EOL> self . input_queue = input_queue <EOL> self . async = async <EOL> if self . async : <EOL> assert self . config . nr_tower > <NUM_LIT:1> <EOL> @ staticmethod <EOL> def _average_grads ( tower_grads ) : <EOL> ret = [ ] <EOL> for grad_and_vars in zip ( * tower_grads ) : <EOL> v = grad_and_vars [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] <EOL> try : <EOL> grad = tf . add_n ( [ x [ <NUM_LIT:0> ] for x in grad_and_vars ] ) / float ( len ( tower_grads ) ) <EOL> except AssertionError : <EOL> logger . error ( "<STR_LIT>" . format ( v . name ) ) <EOL> raise <EOL> ret . append ( ( grad , v ) ) <EOL> return ret <EOL> def _get_model_inputs ( self ) : <EOL> """<STR_LIT>""" <EOL> ret = self . input_queue . dequeue ( ) <EOL> if isinstance ( ret , tf . Tensor ) : <EOL> ret = [ ret ] <EOL> assert len ( ret ) == len ( self . input_vars ) <EOL> for qv , v in zip ( ret , self . input_vars ) : <EOL> qv . set_shape ( v . get_shape ( ) ) <EOL> return ret <EOL> def _single_tower_grad ( self ) : <EOL> """<STR_LIT>""" <EOL> model_inputs = self . _get_model_inputs ( ) <EOL> cost_var = self . model . get_cost ( model_inputs , is_training = True ) <EOL> grads = self . config . optimizer . compute_gradients ( cost_var ) <EOL> tf . add_to_collection ( MOVING_SUMMARY_VARS_KEY , cost_var ) <EOL> return grads <EOL> def _multi_tower_grads ( self ) : <EOL> logger . info ( "<STR_LIT>" . format ( self . config . nr_tower ) ) <EOL> collect_dedup = [ tf . GraphKeys . SUMMARIES , MOVING_SUMMARY_VARS_KEY ] <EOL> kept_summaries = { } <EOL> grad_list = [ ] <EOL> for i in range ( self . config . nr_tower ) : <EOL> with tf . device ( '<STR_LIT>' . format ( i ) ) , tf . name_scope ( '<STR_LIT>' . format ( i ) ) as scope : <EOL> logger . info ( "<STR_LIT>" . format ( i ) ) <EOL> model_inputs = self . _get_model_inputs ( ) <EOL> cost_var = self . model . get_cost ( model_inputs , is_training = True ) <EOL> grad_list . append ( <EOL> self . config . optimizer . compute_gradients ( cost_var , gate_gradients = <NUM_LIT:0> ) ) <EOL> if i == <NUM_LIT:0> : <EOL> tf . add_to_collection ( MOVING_SUMMARY_VARS_KEY , cost_var ) <EOL> tf . get_variable_scope ( ) . reuse_variables ( ) <EOL> for k in collect_dedup : <EOL> kept_summaries [ k ] = copy . copy ( tf . get_collection ( k ) ) <EOL> for k in collect_dedup : <EOL> del tf . get_collection_ref ( k ) [ : ] <EOL> tf . get_collection_ref ( k ) . extend ( kept_summaries [ k ] ) <EOL> return grad_list <EOL> def train ( self ) : <EOL> enqueue_op = self . input_queue . enqueue ( self . input_vars ) <EOL> if self . config . nr_tower > <NUM_LIT:1> : <EOL> grad_list = self . _multi_tower_grads ( ) <EOL> if not self . async : <EOL> grads = QueueInputTrainer . _average_grads ( grad_list ) <EOL> grads = self . process_grads ( grads ) <EOL> else : <EOL> grad_list = [ self . process_grads ( g ) for g in grad_list ] <EOL> def scale ( grads ) : <EOL> return [ ( grad / self . config . nr_tower , var ) for grad , var in grads ] <EOL> grad_list = map ( scale , grad_list ) <EOL> grads = grad_list [ <NUM_LIT:0> ] <EOL> else : <EOL> grads = self . _single_tower_grad ( ) <EOL> grads = self . process_grads ( grads ) <EOL> self . train_op = tf . group ( <EOL> self . config . optimizer . apply_gradients ( grads , get_global_step_var ( ) ) , <EOL> summary_moving_average ( ) ) <EOL> if self . async : <EOL> self . threads = [ ] <EOL> for k in range ( <NUM_LIT:1> , self . config . nr_tower ) : <EOL> train_op = self . config . optimizer . apply_gradients ( grad_list [ k ] ) <EOL> f = lambda op = train_op : self . sess . run ( [ op ] ) <EOL> th = LoopThread ( f ) <EOL> th . pause ( ) <EOL> th . start ( ) <EOL> self . threads . append ( th ) <EOL> self . async_running = False <EOL> self . init_session_and_coord ( ) <EOL> self . input_th = EnqueueThread ( self , self . input_queue , enqueue_op , self . input_vars ) <EOL> self . main_loop ( ) <EOL> def _start_all_threads ( self ) : <EOL> super ( QueueInputTrainer , self ) . _start_all_threads ( ) <EOL> self . input_th . start ( ) <EOL> def run_step ( self ) : <EOL> if self . async : <EOL> if not self . async_running : <EOL> self . async_running = True <EOL> for th in self . threads : <EOL> th . resume ( ) <EOL> self . sess . run ( [ self . train_op ] ) <EOL> def _trigger_epoch ( self ) : <EOL> if self . async : <EOL> self . async_running = False <EOL> for th in self . threads : <EOL> th . pause ( ) <EOL> if self . summary_op is not None : <EOL> summary_str = self . summary_op . eval ( ) <EOL> self . _process_summary ( summary_str ) <EOL> def start_train ( config ) : <EOL> tr = QueueInputTrainer ( config ) <EOL> tr . train ( ) </s>
