deepmd.tf.utils.parallel_op

Module Contents

Classes

ParallelOp

Run an op with data parallelism.
class deepmd.tf.utils.parallel_op.ParallelOp(builder: Callable[Ellipsis, Tuple[Dict[str, deepmd.tf.env.tf.Tensor], Tuple[deepmd.tf.env.tf.Tensor]]], nthreads: int | None = None, config: deepmd.tf.env.tf.ConfigProto | None = None)[source]

Run an op with data parallelism.

Parameters:
builderCallable[…, Tuple[Dict[str, tf.Tensor], Tuple[tf.Tensor]]]

returns two objects: a dict which stores placeholders by key, and a tuple with the final op(s)

nthreadsint, optional

the number of threads

configtf.ConfigProto, optional

tf.ConfigProto

Examples

>>> from deepmd.tf.env import tf
>>> from deepmd.tf.utils.parallel_op import ParallelOp
>>> def builder():
...     x = tf.placeholder(tf.int32, [1])
...     return {"x": x}, (x + 1)
>>> p = ParallelOp(builder, nthreads=4)
>>> def feed():
...     for ii in range(10):
...         yield {"x": [ii]}
>>> print(*p.generate(tf.Session(), feed()))
[1] [2] [3] [4] [5] [6] [7] [8] [9] [10]
generate(sess: deepmd.tf.env.tf.Session, feed: Generator[Dict[str, Any], None, None]) → Generator[Tuple, None, None][source]

Returns a generator.

Parameters:
sesstf.Session

TensorFlow session

feedGenerator[dict, None, None]

generator which yields feed_dict

Yields:
Generator[Tuple, None, None]

generator which yields session returns