Graph Item

GraphItem and its supporting functionality.

cached_property(fn, *args, **kwargs)

Decorator to make a function a “cached property”.

This means that it is a property whose return value is cached after the first time it is called.

Parameters

• fn – The function to be made a cached property

• *args – Any args for the function

• **kwargs – Any kwargs for the function

Returns

function

get_default_graph_item()

Get the current default graph_item under the graph_item scope.

Returns

GraphItem

wrap_optimizer_init(fn: Callable)

Wraps the __init__ function of OptimizerV2 objects and stores the info in the default GraphItem.

wrap_optimizer_apply_gradient(fn: Callable)

Wraps the apply_gradients function of OptimizerV2 objects and stores the info in the default GraphItem.

class Info

Bases: object

Stores useful variable tracking information.

In essence, replaces collections, and this way we don’t have to deal with MetaGraphs.

property initializers

Initializers.

property trainable_variables

Trainable Variables.

property untrainable_variables

Untrainable Variables.

pop_variable(var_name)

Pop out a variable by its name from info tracker.

update_variables(variables, replace=True)

Update variables in GraphItem Info.

Parameters

• variables (Iterable[VariableType]) – Iterable of variables to insert.

• replace (bool) – Whether or not to overwrite existing contents.

update_table_initializers(table_initializers, replace=True)

Update table initializers in GraphItem Info.

Parameters

• table_initializers (Iterable(Union[ops.Operation, str]) – Initializers for lookup tables.

• replace (bool) – Whether or not to overwrite existing contents.

update_savers(savers, replace=True)

Update savers in GraphItem Info.

Parameters

• savers (Iterable[SaverType]) – Iterable of saverdefs to insert.

• replace – Whether or not to overwrite existing contents.

copy()

Copy info.

class GraphItem(graph: Optional[tensorflow.python.framework.ops.Graph] = None, graph_def: Optional[tensorflow.core.framework.graph_pb2.GraphDef] = None)

Bases: object

GraphItem is a TensorFlow Graph wrapper.

It represents the states in-between consecutive AutoDist.kernel graph transformations. tf.Graph is the primary property of GraphItem, whereas MetaGraph is exported/generated on demand.

A GraphItem can be constructed with either a tf.Graph or a GraphDef.

set_optimize()

Start a loop of synchronizer apply.

reset_optimize()

End a loop of synchronizer apply.

get_trainable_variables()

Get variables that need to be synchronized if doing data parallelism.

get_all_variables()

Get all variables in this graph item.

as_default()

A context scope with current graph item as the default.

extend_optimizer_info(optimizer, *args, **kwargs)

Add the detected optimizer to the object.

extend_gradient_info(grads, targets)

Add the detected grad-target pairs to the object.

extend_gradient_info_by_names(grads, targets)

Add the detected grad-target pairs to the object by names.

pop_gradient_info(var_name: str)

Pop out a grad target pair by variable name.

copy_gradient_info_from(other)

Copy gradient info from the another GraphItem object.

property graph

Returns the Graph associated with this GraphItem.

Returns

ops.Graph

property all_update_ops

Get all ops in the graph that perform stateful operations.

Returns

List

property var_op_name_to_grad_info

0”) to its (grad, var, update_op) tuple.

Type

A mapping from VarHandleOp name (e.g. “W” not “W

property var_op_name_to_grad_info_v2

An optimized version that is aware of this method is iteratively used. It optimize based on.

1. Give an updated option, if the graph has not been updated before this query, then it will not calculate again. A new method considering this v2 method needs to manipulate updated outside.

2. Give an var_queried option, which will record which variable has been doen synchronized. If all the variable associated with an update op has been synchronized, this method will not consier the update op next time (it will reconsider if the current loop has done processed, so the set/reset optimize method is necessary to set the boolean flags). This optimization is inspired by that the for loop in this method is executed for every update_op, which is typically a lot, and causes the slowness. This option is safe in that if the var_queried is not set outside, it will not trigger the remove op.

property grad_target_name_pairs

List of names of grad and target variable pairs.

Returns

List

property grad_target_pairs

List of grad and target variable pairs.

Returns

List

property trainable_var_op_to_var

Mapping from trainable variable ops (e.g. VarHandleOps) to the Variables.

Returns

Dict

get_colocation_op(colocation_group)

Get the binding op for a given colocation group.

Parameters

colocation_group (bytes) – The colocation group

Returns

Op

get_ops_in_graph(op_iter)

Given an iterator of ops or op names, return the corresponding ops in self graph.

Parameters

op_iter (Iterable[Union[ops.Operation, str]]) – Ops or ops names

Returns

Iterable

prepare()

Prepare for building strategy and/or transforming graph.

serialize(path)

Serialize a graph_item to a specific proto string down to a file path.

classmethod deserialize(path)

Deserialize a graph_item serialized proto message from a file path.