Skip to content

Commit ffc1deb

Browse files
authored
Merge pull request NVlabs#9 from xivh/master
fixed encoder
2 parents 358a046 + 938ac88 commit ffc1deb

2 files changed

Lines changed: 12 additions & 4 deletions

File tree

dnnlib/tflib/network.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,7 @@ def run(self,
389389
minibatch_size: int = None,
390390
num_gpus: int = 1,
391391
assume_frozen: bool = False,
392+
custom_inputs: Any = None,
392393
**dynamic_kwargs) -> Union[np.ndarray, Tuple[np.ndarray, ...], List[np.ndarray]]:
393394
"""Run this network for the given NumPy array(s), and return the output(s) as NumPy array(s).
394395
@@ -404,6 +405,7 @@ def run(self,
404405
minibatch_size: Maximum minibatch size to use, None = disable batching.
405406
num_gpus: Number of GPUs to use.
406407
assume_frozen: Improve multi-GPU performance by assuming that the trainable parameters will remain changed between calls.
408+
custom_inputs: Allow to use another tensor as input instead of default placeholders.
407409
dynamic_kwargs: Additional keyword arguments to be passed into the network build function.
408410
"""
409411
assert len(in_arrays) == self.num_inputs
@@ -427,10 +429,15 @@ def unwind_key(obj):
427429

428430
# Build graph.
429431
if key not in self._run_cache:
430-
with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
431-
with tf.device("/cpu:0"):
432-
in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
433-
in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
432+
with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
433+
if custom_inputs is not None:
434+
with tf.device("/gpu:0"):
435+
in_expr = [input_builder(name) for input_builder, name in zip(custom_inputs, self.input_names)]
436+
in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
437+
else:
438+
with tf.device("/cpu:0"):
439+
in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
440+
in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
434441

435442
out_split = []
436443
for gpu in range(num_gpus):

encode_images.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import pretrained_networks
1111
from encoder.generator_model import Generator
1212
from encoder.perceptual_model import PerceptualModel
13+
from encoder.perceptual_model import load_images
1314
from keras.models import load_model
1415
from keras.applications.resnet50 import preprocess_input
1516

0 commit comments

Comments
 (0)