Commit 2aa760b1 authored by Yuxin Wu's avatar Yuxin Wu

fix shape comparison when reusing placeholders. fix #1329

parent eb25cd7f
......@@ -6,7 +6,7 @@
There are two ways to do inference during training.
1. The easiest way is to write a callback, and use
[self.trainer.get_predictor()](../modules/modules/train.html#tensorpack.train.TowerTrainer.get_predictor)
[self.trainer.get_predictor()](../modules/train.html#tensorpack.train.TowerTrainer.get_predictor)
to get a callable under inference mode.
See [Write a Callback](extend/callback.html).
......
......@@ -34,9 +34,12 @@ def build_or_reuse_placeholder(tensor_spec):
assert "Placeholder" in tensor.op.type, "Tensor {} exists but is not a placeholder!".format(name)
assert tensor_spec.is_compatible_with(tensor), \
"Tensor {} exists but is not compatible with the signature!".format(tensor)
if tensor.shape == tensor_spec.shape:
if tensor.shape.as_list() == tensor_spec.shape.as_list():
# It might be desirable to use a placeholder of a different shape in some tower
# (e.g., a less specific shape)
# Comparing `tensor.shape` directly doesn't work, because
# tensorflow thinks `tf.Dimension(None)` and `tf.Dimension(None)` are not equal.
return tensor
except KeyError:
pass
......
......@@ -46,7 +46,8 @@ class TowerTrainer(Trainer):
def tower_func(self):
"""
A :class:`TowerFunc` instance.
See [tutorial on tower function](http://tensorpack.readthedocs.io/tutorial/trainer.html#tower-trainer)
See `tutorial on tower function
<http://tensorpack.readthedocs.io/tutorial/trainer.html#tower-trainer>`_
for more information.
"""
return self._tower_func
......
......@@ -388,8 +388,6 @@ class HorovodTrainer(SingleCostTrainer):
certain numerical issues in practice.
"""
BROADCAST_EVERY_EPOCH = True
def __init__(self, average=True, compression=None):
"""
Args:
......@@ -415,6 +413,8 @@ class HorovodTrainer(SingleCostTrainer):
logger.info("[HorovodTrainer] local rank={}".format(self._local_rank))
super(HorovodTrainer, self).__init__()
self.BROADCAST_EVERY_EPOCH = True
def mpi_enabled(self):
"""
Returns:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment