class FeedfreePredictor(PredictorBase):
"""
Create a predictor that takes inputs from an :class:`InputSource`, instead of from feeds.
An instance `pred` of :class:`FeedfreePredictor` can be called only by `pred()`, which returns
a list of output values as defined in config.output_names.
"""
def __init__(self, config, input_source):
"""
Args:
config (PredictConfig): the config to use.
input_source (InputSource): the feedfree InputSource to use.
Must match the inputs_desc in config.
"""
self._config = config
self._input_source = input_source
assert config.return_input is False, \
"return_input is not supported in FeedfreePredictor! " \
"If you need to fetch inputs, add the names to the output_names!"
self._hooks = []
self.graph = config._maybe_create_graph()
with self.graph.as_default():
self._input_callbacks = Callbacks(
self._input_source.setup(config.inputs_desc))
with PredictTowerContext(''):
self._input_tensors = self._input_source.get_input_tensors()
config.tower_func(*self._input_tensors)
self._tower_handle = config.tower_func.towers[-1]
self._output_tensors = self._tower_handle.get_tensors(config.output_names)
self._input_callbacks.setup_graph(None)
for h in self._input_callbacks.get_hooks():
self._register_hook(h)
self._initialize_session()
def _register_hook(self, hook):
"""
Args:
hook (tf.train.SessionRunHook):
"""
self._hooks.append(hook)
def _initialize_session(self):
# init the session
self._config.session_init._setup_graph()
self._sess = self._config.session_creator.create_session()
self._config.session_init._run_init(self._sess)
with self._sess.as_default():
self._input_callbacks.before_train()
self._hooked_sess = HookedSession(self._sess, self._hooks)
def __call__(self):
return self._hooked_sess.run(self._output_tensors)
def _do_call(self):
raise NotImplementedError("You're calling the wrong function!")
def _initialize_session(self):
# init the session
self._config.session_init._setup_graph()
self._sess = self._config.session_creator.create_session()
self._config.session_init._run_init(self._sess)
with self._sess.as_default():
self._input_callbacks.before_train()
self._hooked_sess = HookedSession(self._sess, self._hooks)
class DataParallelInferenceRunner(InferenceRunnerBase):
"""
Inference with data-parallel support on multiple GPUs.
It will build one predict tower on each GPU, and run prediction
with a large total batch in parallel on all GPUs.
It will run the remainder (when the total size of input is not a multiple of #GPU)
sequentially.
"""
def __init__(self,
input,
infs,
gpus,
tower_name='InferenceTower',
tower_func=None):
"""
Args:
input (DataFlow or QueueInput)
gpus (int or list[int]): #gpus, or list of GPU id
tower_name (str): the name scope of the tower to build. Need to set a
different one if multiple InferenceRunner are used.
tower_func (tfutils.TowerFuncWrapper or None): the tower function to be used to build the graph.
The tower function will be called under a `training=False` TowerContext.
The default is `trainer.tower_func`,
but you can change it to a different tower function
if you need to inference with several different models.
"""
if isinstance(gpus, int):
gpus = list(range(gpus))
self._devices = [_device_from_int(k) for k in gpus]
self._tower_names = [
'{}{}'.format(tower_name, k) for k in range(len(gpus))
]
if isinstance(input, DataFlow):
input = QueueInput(input)
assert isinstance(input, QueueInput), input
super(DataParallelInferenceRunner, self).__init__(input, infs)
assert self._size > 0, "Input for DataParallelInferenceRunner must have a size!"
self._hooks = []
self._hooks_parallel = []
self._tower_func = tower_func
def _setup_graph(self):
self._handles = []
if self._tower_func is None:
assert self.trainer.tower_func is not None, "You must set tower_func of the trainer to use InferenceRunner!"
self._tower_func = self.trainer.tower_func
input_callbacks = self._input_source.setup(
self._tower_func.inputs_desc)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for idx, dev in enumerate(self._devices):
vs_name = self.trainer._vs_name_for_predictor(idx)
with tf.device(dev), PredictTowerContext(
self._tower_names[idx], vs_name=vs_name):
logger.info(
"[InferenceRunner] Building tower '{}' on device {} {}..."
.format(
self._tower_names[idx], dev,
"with variable scope '{}'".format(vs_name)
if vs_name else ''))
# TODO log for tower creation, here or in tower.py?
self._tower_func(*self._input_source.get_input_tensors())
self._handles.append(self._tower_func.towers[-1])
# setup callbacks and hooks
self._input_callbacks = Callbacks(input_callbacks)
# TODO InputSource might have hooks which break us.
# e.g. hooks from StagingInput will force the consumption
# of nr_tower datapoints in every run.
input_hooks = self._input_callbacks.get_hooks()
self._hooks.extend([self._build_hook(inf)
for inf in self.infs] + input_hooks)
self._hooks_parallel.extend(
[self._build_hook_parallel(inf)
for inf in self.infs] + input_hooks)
for inf in self.infs:
inf.setup_graph(self.trainer)
self._input_callbacks.setup_graph(self.trainer)
def register_hook(self, h):
logger.info(
"[DataParallelInferenceRunner] Registering hook {} on both parallel and sequential inference."
)
self._hooks.append(h)
self._hooks_parallel.append(h)
class _InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
"""
Args:
size(int): number of tensors to fetch per tower
"""
super(DataParallelInferenceRunner._InferencerToHookDataParallel,
self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.on_fetches(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_fetches()
sz = len(out_names)
fetches = list(
itertools.chain(*[t.get_tensors(out_names)
for t in self._handles]))
return self._InferencerToHookDataParallel(inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_fetches()
fetches = self._handles[0].get_tensors(out_names)
return InferencerToHook(inf, fetches)
def _before_train(self):
super(DataParallelInferenceRunner, self)._before_train()
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_epoch()
total = self._size
nr_tower = len(self._devices)
self._input_source.reset_state()
with _inference_context():
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
self._parallel_hooked_sess.run(fetches=[])
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
for _ in range(total):
self._hooked_sess.run(fetches=[])
pbar.update(1)
for inf in self.infs:
inf.trigger_epoch()
class DataParallelInferenceRunner(InferenceRunnerBase):
"""
Inference with data-parallel support on multiple GPUs.
It will build one predict tower on each GPU, and run prediction
with a large total batch in parallel on all GPUs.
It will run the remainder (when the total size of input is not a multiple of #GPU)
sequentially.
"""
def __init__(self, input, infs, gpus):
"""
Args:
input (DataFlow or QueueInput)
gpus (int or list[int]): #gpus, or list of GPU id
"""
if isinstance(gpus, int):
gpus = list(range(gpus))
self._tower_names = ['InferenceTower{}'.format(k) for k in range(len(gpus))]
if isinstance(input, DataFlow):
input = QueueInput(input)
assert isinstance(input, QueueInput), input
super(DataParallelInferenceRunner, self).__init__(input, infs)
assert self._size > 0, "Input for DataParallelInferenceRunner must have a size!"
self._gpus = gpus
self._hooks = []
self._hooks_parallel = []
def _setup_graph(self):
self._handles = []
assert self.trainer.tower_func is not None, "You must set tower_func of the trainer to use InferenceRunner!"
input_callbacks = self._input_source.setup(self.trainer.inputs_desc)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for idx, t in enumerate(self._gpus):
tower_name = self._tower_names[idx]
SimplePredictBuilder(
ns_name=tower_name,
vs_name=self.trainer._main_tower_vs_name, device=t).build(
self._input_source, self.trainer.tower_func)
self._handles.append(self.trainer.tower_func.towers[-1])
# setup callbacks and hooks
self._input_callbacks = Callbacks(input_callbacks)
# TODO InputSource might have hooks which break us.
# e.g. hooks from StagingInput will force the consumption
# of nr_tower datapoints in every run.
input_hooks = self._input_callbacks.get_hooks()
self._hooks.extend([self._build_hook(inf) for inf in self.infs] + input_hooks)
self._hooks_parallel.extend([self._build_hook_parallel(inf) for inf in self.infs] + input_hooks)
for inf in self.infs:
inf.setup_graph(self.trainer)
self._input_callbacks.setup_graph(self.trainer)
def register_hook(self, h):
logger.info(
"[DataParallelInferenceRunner] Registering hook {} on both parallel and sequential inference.")
self._hooks.append(h)
self._hooks_parallel.append(h)
class InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
"""
Args:
size(int): number of tensors to fetch per tower
"""
super(DataParallelInferenceRunner.InferencerToHookDataParallel, self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.on_fetches(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_fetches()
sz = len(out_names)
fetches = list(itertools.chain(*[t.get_tensors(out_names) for t in self._handles]))
return self.InferencerToHookDataParallel(inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_fetches()
fetches = self._handles[0].get_tensors(out_names)
return InferencerToHook(inf, fetches)
def _before_train(self):
super(DataParallelInferenceRunner, self)._before_train()
self._parallel_hooked_sess = HookedSession(self.trainer.sess, self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_epoch()
total = self._size
nr_tower = len(self._gpus)
with _inference_context():
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
self._parallel_hooked_sess.run(fetches=[])
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
for _ in range(total):
self._hooked_sess.run(fetches=[])
pbar.update(1)
for inf in self.infs:
inf.trigger_epoch()
def _before_train(self):
super(DataParallelInferenceRunner, self)._before_train()
self._parallel_hooked_sess = HookedSession(self.trainer.sess, self._hooks_parallel)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
class InferenceRunnerBase(Callback):
""" Base methods for inference runner"""
def __init__(self,
input,
infs,
input_names=None,
prefix='',
extra_hooks=None):
"""
Args:
input (InputSource): the input to use. Must have ``size()``.
infs (list): list of :class:`Inferencer` to run.
input_names (list): must be a subset of the names in InputDesc.
prefix(str): an prefix used to build the tower. Must be set
differently if more than one :class:`InferenceRunner` are used.
extra_hooks (list): extra ``SessionRunHook`` to run with the evaluation.
"""
self._input_source = input
if not isinstance(infs, list):
self.infs = [infs]
else:
self.infs = infs
for v in self.infs:
assert isinstance(v, Inferencer), v
if input_names is not None:
assert isinstance(input_names, list)
self.input_names = input_names
try:
self._size = input.size()
except NotImplementedError:
raise ValueError("Input used in InferenceRunner must have a size!")
self._prefix = prefix
if extra_hooks is None:
extra_hooks = []
self._extra_hooks = extra_hooks
def _setup_input_names(self):
# just use all the placeholders, if input_name is None
if self.input_names is None:
inputs = self.trainer.model.get_reused_placehdrs()
self.input_names = [x.name for x in inputs]
# TODO sparse. even if it works here, sparse still is unavailable
# because get_tensor_by_name doesn't work for sparse
# def get_name(x):
# if isinstance(x, tf.SparseTensor):
# return x.op.name.split('/')[0]
# return x.name
def _setup_graph(self):
self._input_source.setup(self.trainer.model)
self._setup_input_names()
# Use predict_tower in train config. either gpuid or -1
self._predict_tower_id = self.trainer.config.predict_tower[0]
in_tensors = self._find_input_tensors()
assert isinstance(in_tensors, list), in_tensors
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
def fn(_):
self.trainer.model.build_graph(in_tensors)
PredictorTowerBuilder(fn,
self._prefix).build(self._predict_tower_id)
self._feed_tensors = self._find_feed_tensors()
self._hooks = [self._build_hook(inf) for inf in self.infs]
def _before_train(self):
self._hooks.extend(self._extra_hooks)
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
def _get_tensors_maybe_in_tower(self, names):
placeholder_names = set(
[k.name for k in self.trainer.model.get_inputs_desc()])
get_tensor_fn = PredictorTowerBuilder.get_tensors_maybe_in_tower
return get_tensor_fn(placeholder_names,
names,
self._predict_tower_id,
prefix=self._prefix)
def _find_input_tensors(self):
pass
@abstractmethod
def _find_feed_tensors(self):
pass
@abstractmethod
def _build_hook(self, inf):
pass
def _trigger(self):
for inf in self.infs:
inf.before_inference()
# iterate over the data, and run the hooked session
self._input_source.reset_state()
for _ in tqdm.trange(self._input_source.size(), **get_tqdm_kwargs()):
dp = self._input_source.next_feed()
feed = dict(zip(self._feed_tensors, dp))
self._hooked_sess.run(fetches=[], feed_dict=feed)
summary_inferencer(self.trainer, self.infs)
def _before_train(self):
self._hooks.extend(self._extra_hooks)
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
class DataParallelInferenceRunner(InferenceRunner):
def __init__(self, ds, infs, gpus, input_names=None):
super(DataParallelInferenceRunner,
self).__init__(ds, infs, input_names)
self._gpus = gpus
def _setup_graph(self):
model = self.trainer.model
self._input_source.setup(model)
self._setup_input_names()
# build graph
def build_tower(k):
towername = TowerContext.get_predict_tower_name(k)
# inputs (placeholders) for this tower only
input_tensors = model.build_placeholders(prefix=towername + '/')
model.build_graph(input_tensors)
builder = PredictorTowerBuilder(build_tower, prefix=self._prefix)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for t in self._gpus:
builder.build(t)
# setup feeds and hooks
self._feed_tensors = self._find_feed_tensors()
self._hooks_parallel = [
self._build_hook_parallel(inf) for inf in self.infs
]
self._hooks = [self._build_hook(inf) for inf in self.infs]
def _duplicate_names_across_towers(self, names):
ret = []
for t in self._gpus:
ret.extend([
TowerContext.get_predict_tower_name(t, self._prefix) + '/' + n
for n in names
])
return ret
def _find_feed_tensors(self):
names = self._duplicate_names_across_towers(self.input_names)
return get_tensors_by_names(names)
class InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
super(DataParallelInferenceRunner.InferencerToHookDataParallel,
self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.datapoint(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_output_tensors()
sz = len(out_names)
out_names = self._duplicate_names_across_towers(out_names)
fetches = get_tensors_by_names(out_names)
return DataParallelInferenceRunner.InferencerToHookDataParallel(
inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_output_tensors()
names = [
TowerContext.get_predict_tower_name(self._gpus[0], self._prefix) +
'/' + n for n in out_names
]
fetches = get_tensors_by_names(names)
return InferencerToHook(inf, fetches)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_inference()
self._input_source.reset_state()
total = self._input_source.size()
nr_tower = len(self._gpus)
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
dps = []
for k in self._gpus:
dps.extend(self._input_source.next_feed())
feed = dict(zip(self._feed_tensors, dps))
self._parallel_hooked_sess.run(fetches=[], feed_dict=feed)
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
while total > 0:
dp = self._input_source.next_feed()
feed = dict(zip(self._feed_tensors[:len(dp)], dp))
self._hooked_sess.run(fetches=[], feed_dict=feed)
summary_inferencer(self.trainer, self.infs)
class DataParallelInferenceRunner(InferenceRunnerBase):
"""
Inference with data-parallel support on multiple GPUs.
It will build one predict tower on each GPU, and run prediction
with a larger batch.
"""
def __init__(self, input, infs, gpus):
"""
Args:
input (DataFlow or QueueInput)
gpus (list[int]): list of GPU id
"""
self._tower_names = [
'InferenceTower{}'.format(k) for k in range(len(gpus))
]
if isinstance(input, DataFlow):
input = QueueInput(input)
assert isinstance(input, QueueInput), input
super(DataParallelInferenceRunner, self).__init__(input, infs)
self._gpus = gpus
def _setup_graph(self):
self._handles = []
assert self.trainer.tower_func is not None, "You must set tower_func of the trainer to use InferenceRunner!"
input_callbacks = self._input_source.setup(self.trainer.inputs_desc)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for idx, t in enumerate(self._gpus):
tower_name = self._tower_names[idx]
SimplePredictBuilder(ns_name=tower_name,
vs_name=self.trainer._main_tower_vs_name,
device=t).build(self._input_source,
self.trainer.tower_func)
self._handles.append(self.trainer.tower_func.towers[-1])
# setup callbacks and hooks
self._input_callbacks = Callbacks(input_callbacks)
# TODO InputSource might have hooks which break us.
# e.g. hooks from StagingInput will force the consumption
# of nr_tower datapoints in every run.
input_hooks = self._input_callbacks.get_hooks()
self._hooks = [self._build_hook(inf)
for inf in self.infs] + input_hooks
self._hooks_parallel = [
self._build_hook_parallel(inf) for inf in self.infs
] + input_hooks
for inf in self.infs:
inf.setup_graph(self.trainer)
self._input_callbacks.setup_graph(self.trainer)
def register_hook(self, h):
raise NotImplementedError(
"DataParallelInferenceRunner doesn't accept extra hooks!")
class InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
"""
Args:
size(int): number of tensors to fetch per tower
"""
super(DataParallelInferenceRunner.InferencerToHookDataParallel,
self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.on_fetches(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_fetches()
sz = len(out_names)
fetches = list(
itertools.chain(*[t.get_tensors(out_names)
for t in self._handles]))
return self.InferencerToHookDataParallel(inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_fetches()
fetches = self._handles[0].get_tensors(out_names)
return InferencerToHook(inf, fetches)
def _before_train(self):
super(DataParallelInferenceRunner, self)._before_train()
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_epoch()
self._input_source.reset_state()
total = self._size
nr_tower = len(self._gpus)
with _inference_context():
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
self._parallel_hooked_sess.run(fetches=[])
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
for _ in range(total):
self._hooked_sess.run(fetches=[])
pbar.update(1)
for inf in self.infs:
inf.trigger_epoch()
class InferenceRunnerBase(Callback):
""" Base methods for inference runner"""
def __init__(self,
input,
infs,
tower_name='InferenceTower',
extra_hooks=None,
prefix=None):
"""
Args:
input (InputSource): the input to use. Must have ``size()``.
infs (list[Inferencer]): list of :class:`Inferencer` to run.
tower_name(str): name scope to build the tower. Must be set
differently if more than one :class:`InferenceRunner` are used.
extra_hooks (list[SessionRunHook]): extra :class:`SessionRunHook` to run with the evaluation.
"""
self._input_source = input
if not isinstance(infs, list):
self.infs = [infs]
else:
self.infs = infs
for v in self.infs:
assert isinstance(v, Inferencer), v
try:
self._size = input.size()
except NotImplementedError:
raise ValueError("Input used in InferenceRunner must have a size!")
self._tower_name = tower_name
if prefix is not None:
self._tower_name = 'InferenceTower' + prefix
if extra_hooks is None:
extra_hooks = []
self._extra_hooks = extra_hooks
def _setup_graph(self):
# Use predict_tower in train config. either gpuid or -1
tower_id = self.trainer.config.predict_tower[0]
device = '/gpu:{}'.format(tower_id) if tower_id >= 0 else '/cpu:0'
self._input_source.setup(self.trainer.model.get_inputs_desc())
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
self._tower_handle = self.trainer.predictor_factory.build(
self._tower_name, device, self._input_source)
self._hooks = [self._build_hook(inf) for inf in self.infs]
cbs = self._input_source.get_callbacks()
self._hooks.extend([CallbackToHook(cb) for cb in cbs])
def _before_train(self):
self._hooks.extend(self._extra_hooks)
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
@abstractmethod
def _build_hook(self, inf):
pass
def _trigger(self):
for inf in self.infs:
inf.before_inference()
# iterate over the data, and run the hooked session
self._input_source.reset_state()
for _ in tqdm.trange(self._size, **get_tqdm_kwargs()):
self._hooked_sess.run(fetches=[])
summary_inferencer(self.trainer, self.infs)
class InferenceRunnerBase(Callback):
""" Base methods for inference runner"""
def __init__(self, input, infs, prefix='', extra_hooks=None):
"""
Args:
input (InputSource): the input to use. Must have ``size()``.
infs (list[Inferencer]): list of :class:`Inferencer` to run.
prefix(str): an prefix used to build the tower. Must be set
differently if more than one :class:`InferenceRunner` are used.
extra_hooks (list[SessionRunHook]): extra :class:`SessionRunHook` to run with the evaluation.
"""
self._input_source = input
if not isinstance(infs, list):
self.infs = [infs]
else:
self.infs = infs
for v in self.infs:
assert isinstance(v, Inferencer), v
try:
self._size = input.size()
except NotImplementedError:
raise ValueError("Input used in InferenceRunner must have a size!")
self._prefix = prefix
if extra_hooks is None:
extra_hooks = []
self._extra_hooks = extra_hooks
def _setup_graph(self):
self._input_source.setup(self.trainer.model)
# Use predict_tower in train config. either gpuid or -1
self._predict_tower_id = self.trainer.config.predict_tower[0]
def fn(_):
in_tensors = self._input_source.get_input_tensors()
self.trainer.model.build_graph(in_tensors)
with tf.variable_scope(self.trainer.vs_name_for_predictor, reuse=True):
PredictorTowerBuilder(fn, self._prefix).build(self._predict_tower_id)
self._hooks = [self._build_hook(inf) for inf in self.infs]
def _before_train(self):
self._hooks.extend(self._extra_hooks)
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
def _get_tensors_maybe_in_tower(self, names):
placeholder_names = set([k.name for k in self.trainer.model.get_inputs_desc()])
get_tensor_fn = PredictorTowerBuilder.get_tensors_maybe_in_tower
return get_tensor_fn(placeholder_names, names, self._predict_tower_id, prefix=self._prefix)
@abstractmethod
def _build_hook(self, inf):
pass
def _trigger(self):
for inf in self.infs:
inf.before_inference()
# iterate over the data, and run the hooked session
self._input_source.reset_state()
for _ in tqdm.trange(self._input_source.size(), **get_tqdm_kwargs()):
feed = self._input_source.next_feed()
self._hooked_sess.run(fetches=[], feed_dict=feed)
summary_inferencer(self.trainer, self.infs)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
class InferenceRunnerBase(Callback):
""" Base methods for inference runner"""
def __init__(self, input, infs, input_names=None, prefix='', extra_hooks=None):
"""
Args:
input (InputData): the input to use. Must have ``size()``.
infs (list): list of :class:`Inferencer` to run.
input_names (list): must be a subset of the names in InputDesc.
prefix(str): an prefix used to build the tower. Must be set
differently if more than one :class:`InferenceRunner` are used.
extra_hooks (list): extra ``SessionRunHook`` to run with the evaluation.
"""
self._input_data = input
if not isinstance(infs, list):
self.infs = [infs]
else:
self.infs = infs
for v in self.infs:
assert isinstance(v, Inferencer), v
if input_names is not None:
assert isinstance(input_names, list)
self.input_names = input_names
try:
self._size = input.size()
except NotImplementedError:
raise ValueError("Input used in InferenceRunner must have a size!")
self._prefix = prefix
if extra_hooks is None:
extra_hooks = []
self._extra_hooks = extra_hooks
def _setup_input_names(self):
# just use all the placeholders, if input_name is None
if self.input_names is None:
inputs = self.trainer.model.get_reused_placehdrs()
self.input_names = [x.name for x in inputs]
# TODO sparse. even if it works here, sparse still is unavailable
# because get_tensor_by_name doesn't work for sparse
# def get_name(x):
# if isinstance(x, tf.SparseTensor):
# return x.op.name.split('/')[0]
# return x.name
def _setup_graph(self):
self._input_data.setup(self.trainer.model)
self._setup_input_names()
in_tensors = self._find_input_tensors()
assert isinstance(in_tensors, list), in_tensors
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
def fn(_):
self.trainer.model.build_graph(in_tensors)
PredictorTowerBuilder(fn, self._prefix).build(0)
self._feed_tensors = self._find_feed_tensors()
self._hooks = [self._build_hook(inf) for inf in self.infs]
def _before_train(self):
self._hooks.extend(self._extra_hooks)
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
def _get_tensors_maybe_in_tower(self, names):
placeholder_names = set([k.name for k in self.trainer.model.get_inputs_desc()])
get_tensor_fn = PredictorTowerBuilder.get_tensors_maybe_in_tower
return get_tensor_fn(placeholder_names, names, 0, prefix=self._prefix)
def _find_input_tensors(self):
pass
@abstractmethod
def _find_feed_tensors(self):
pass
@abstractmethod
def _build_hook(self, inf):
pass
def _trigger(self):
for inf in self.infs:
inf.before_inference()
# iterate over the data, and run the hooked session
self._input_data.reset_state()
for _ in tqdm.trange(self._input_data.size(), **get_tqdm_kwargs()):
dp = self._input_data.next_feed()
feed = dict(zip(self._feed_tensors, dp))
self._hooked_sess.run(fetches=[], feed_dict=feed)
summary_inferencer(self.trainer, self.infs)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
self._parallel_hooked_sess = HookedSession(self.trainer.sess, self._hooks_parallel)
class DataParallelInferenceRunner(InferenceRunner):
def __init__(self, ds, infs, gpus, input_names=None):
super(DataParallelInferenceRunner, self).__init__(ds, infs, input_names)
self._gpus = gpus
def _setup_graph(self):
model = self.trainer.model
self._input_data.setup(model)
self._setup_input_names()
# build graph
def build_tower(k):
towername = TowerContext.get_predict_tower_name(k)
# inputs (placeholders) for this tower only
input_tensors = model.build_placeholders(
prefix=towername + '/')
model.build_graph(input_tensors)
builder = PredictorTowerBuilder(build_tower, prefix=self._prefix)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for t in self._gpus:
builder.build(t)
# setup feeds and hooks
self._feed_tensors = self._find_feed_tensors()
self._hooks_parallel = [self._build_hook_parallel(inf) for inf in self.infs]
self._hooks = [self._build_hook(inf) for inf in self.infs]
def _duplicate_names_across_towers(self, names):
ret = []
for t in self._gpus:
ret.extend([TowerContext.get_predict_tower_name(t, self._prefix) +
'/' + n for n in names])
return ret
def _find_feed_tensors(self):
names = self._duplicate_names_across_towers(self.input_names)
return get_tensors_by_names(names)
class InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
super(DataParallelInferenceRunner.InferencerToHookDataParallel, self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.datapoint(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_output_tensors()
sz = len(out_names)
out_names = self._duplicate_names_across_towers(out_names)
fetches = get_tensors_by_names(out_names)
return DataParallelInferenceRunner.InferencerToHookDataParallel(
inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_output_tensors()
names = [TowerContext.get_predict_tower_name(
self._gpus[0], self._prefix) + '/' + n for n in out_names]
fetches = get_tensors_by_names(names)
return InferencerToHook(inf, fetches)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
self._parallel_hooked_sess = HookedSession(self.trainer.sess, self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_inference()
self._input_data.reset_state()
total = self._input_data.size()
nr_tower = len(self._gpus)
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
dps = []
for k in self._gpus:
dps.extend(self._input_data.next_feed())
feed = dict(zip(self._feed_tensors, dps))
self._parallel_hooked_sess.run(fetches=[], feed_dict=feed)
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
while total > 0:
dp = self._input_data.next_feed()
feed = dict(zip(self._feed_tensors[:len(dp)], dp))
self._hooked_sess.run(fetches=[], feed_dict=feed)
summary_inferencer(self.trainer, self.infs)
def _before_train(self):
self._hooks.extend(self._extra_hooks)
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
class DataParallelInferenceRunner(InferenceRunnerBase):
"""
Inference by feeding datapoints in a data-parallel way to multiple GPUs.
Doesn't support remapped InputSource for now.
"""
def __init__(self, input, infs, gpus):
"""
Args:
input (DataParallelFeedInput or DataFlow)
gpus (list[int]): list of GPU id
"""
if isinstance(input, DataFlow):
tower_names = [
TowerContext.get_predict_tower_name(k)
for k in range(len(gpus))
]
input = DataParallelFeedInput(input, tower_names)
assert isinstance(input, DataParallelFeedInput), input
super(DataParallelInferenceRunner, self).__init__(input, infs)
self._gpus = gpus
def _setup_graph(self):
model = self.trainer.model
self._input_source.setup(model.get_inputs_desc())
# build graph
def build_tower(k):
# inputs (placeholders) for this tower only
model.build_graph(self._input_source)
builder = PredictorTowerBuilder(build_tower, prefix=self._prefix)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for t in self._gpus:
builder.build(t)
# setup feeds and hooks
self._hooks_parallel = [
self._build_hook_parallel(inf) for inf in self.infs
]
self._hooks = [self._build_hook(inf) for inf in self.infs]
cbs = self._input_source.get_callbacks()
self._hooks_parallel.extend([CallbackToHook(cb) for cb in cbs])
def _duplicate_names_across_towers(self, names):
ret = []
for t in self._gpus:
ret.extend([
TowerContext.get_predict_tower_name(t, self._prefix) + '/' + n
for n in names
])
return ret
class InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
super(DataParallelInferenceRunner.InferencerToHookDataParallel,
self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.datapoint(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_output_tensors()
sz = len(out_names)
out_names = self._duplicate_names_across_towers(out_names)
fetches = get_tensors_by_names(out_names)
return DataParallelInferenceRunner.InferencerToHookDataParallel(
inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_output_tensors()
names = [
TowerContext.get_predict_tower_name(self._gpus[0], self._prefix) +
'/' + n for n in out_names
]
fetches = get_tensors_by_names(names)
return InferencerToHook(inf, fetches)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_inference()
self._input_source.reset_state()
total = self._size
nr_tower = len(self._gpus)
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
self._parallel_hooked_sess.run(fetches=[])
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
try:
while total > 0:
# TODO XXX doesn't support remap
feed = self._input_source.next_feed(cnt=1)
self._hooked_sess.run(fetches=[], feed_dict=feed)
pbar.update(1)
total -= 1
except AttributeError:
logger.error(
"[DataParallelInferenceRunner] doesn't support InputSource wrappers very well!"
)
logger.error(
"[DataParallelInferenceRunner] Skipping the rest of the datapoints ..."
)
summary_inferencer(self.trainer, self.infs)
def _before_train(self):
super(DataParallelInferenceRunner, self)._before_train()
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
class DataParallelInferenceRunner(InferenceRunnerBase):
"""
Inference by feeding datapoints in a data-parallel way to multiple GPUs.
Doesn't support remapped InputSource for now.
"""
def __init__(self, input, infs, gpus):
"""
Args:
input (DataParallelFeedInput or DataFlow)
gpus (list[int]): list of GPU id
"""
self._tower_names = [
'InferenceTower{}'.format(k) for k in range(len(gpus))
]
if isinstance(input, DataFlow):
input = DataParallelFeedInput(input, self._tower_names)
assert isinstance(input, DataParallelFeedInput), input
super(DataParallelInferenceRunner, self).__init__(input, infs)
self._gpus = gpus
def _setup_graph(self):
self._input_source.setup(self.trainer.model.get_inputs_desc())
self._handles = []
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
for idx, t in enumerate(self._gpus):
tower_name = self._tower_names[idx]
device = '/gpu:{}'.format(t)
self._handles.append(
self.trainer.predictor_factory.build(
tower_name, device, self._input_source))
# setup feeds and hooks
self._hooks_parallel = [
self._build_hook_parallel(inf) for inf in self.infs
]
self._hooks = [self._build_hook(inf) for inf in self.infs]
cbs = self._input_source.get_callbacks()
self._hooks_parallel.extend([CallbackToHook(cb) for cb in cbs])
class InferencerToHookDataParallel(InferencerToHook):
def __init__(self, inf, fetches, size):
"""
Args:
size(int): number of tensors to fetch per tower
"""
super(DataParallelInferenceRunner.InferencerToHookDataParallel,
self).__init__(inf, fetches)
assert len(self._fetches) % size == 0
self._sz = size
def after_run(self, _, run_values):
res = run_values.results
for i in range(0, len(res), self._sz):
vals = res[i:i + self._sz]
self._inf.datapoint(vals)
def _build_hook_parallel(self, inf):
out_names = inf.get_output_tensors()
sz = len(out_names)
fetches = list(
itertools.chain(*[t.get_tensors(out_names)
for t in self._handles]))
return self.InferencerToHookDataParallel(inf, fetches, sz)
def _build_hook(self, inf):
out_names = inf.get_output_tensors()
fetches = self._handles[0].get_tensors(out_names)
return InferencerToHook(inf, fetches)
def _before_train(self):
self._hooked_sess = HookedSession(self.trainer.sess, self._hooks)
self._parallel_hooked_sess = HookedSession(self.trainer.sess,
self._hooks_parallel)
def _trigger(self):
for inf in self.infs:
inf.before_inference()
self._input_source.reset_state()
total = self._size
nr_tower = len(self._gpus)
with tqdm.tqdm(total=total, **get_tqdm_kwargs()) as pbar:
while total >= nr_tower:
self._parallel_hooked_sess.run(fetches=[])
pbar.update(nr_tower)
total -= nr_tower
# take care of the rest
try:
while total > 0:
# TODO XXX doesn't support remap
feed = self._input_source.next_feed(cnt=1)
self._hooked_sess.run(fetches=[], feed_dict=feed)
pbar.update(1)
total -= 1
except AttributeError:
logger.error(
"[DataParallelInferenceRunner] doesn't support InputSource wrappers very well!"
)
logger.error(
"[DataParallelInferenceRunner] Skipping the rest of the datapoints ..."
)
summary_inferencer(self.trainer, self.infs)
