Adds a collection to MetaGraphDef protocol buffer. Args: meta_graph_def: MetaGraphDef protocol buffer. key: One of the GraphKeys or user-defined string. graph: The `Graph` from which to get collections. export_scope: Optional `string`. Name scope to remove. exclude_nodes: An i
(meta_graph_def, key, graph=None,
export_scope=None, exclude_nodes=None,
override_contents=None)
| 378 | |
| 379 | |
| 380 | def add_collection_def(meta_graph_def, key, graph=None, |
| 381 | export_scope=None, exclude_nodes=None, |
| 382 | override_contents=None): |
| 383 | """Adds a collection to MetaGraphDef protocol buffer. |
| 384 | |
| 385 | Args: |
| 386 | meta_graph_def: MetaGraphDef protocol buffer. |
| 387 | key: One of the GraphKeys or user-defined string. |
| 388 | graph: The `Graph` from which to get collections. |
| 389 | export_scope: Optional `string`. Name scope to remove. |
| 390 | exclude_nodes: An iterable of nodes or `string` node names to omit from the |
| 391 | collection, or None. |
| 392 | override_contents: An iterable of values to place in the collection, |
| 393 | ignoring the current values (if set). |
| 394 | """ |
| 395 | if graph and not isinstance(graph, ops.Graph): |
| 396 | raise TypeError("graph must be of type Graph, not %s", type(graph)) |
| 397 | |
| 398 | if not isinstance(key, six.string_types) and not isinstance(key, bytes): |
| 399 | logging.warning("Only collections with string type keys will be " |
| 400 | "serialized. This key has %s", type(key)) |
| 401 | return |
| 402 | |
| 403 | # Sets graph to default graph if it's not passed in. |
| 404 | graph = graph or ops.get_default_graph() |
| 405 | |
| 406 | if override_contents: |
| 407 | collection_list = override_contents |
| 408 | else: |
| 409 | collection_list = graph.get_collection(key) |
| 410 | |
| 411 | # Remove nodes that should not be exported from the collection list. |
| 412 | collection_list = [x for x in collection_list if |
| 413 | _should_include_node(x, export_scope, exclude_nodes)] |
| 414 | if not collection_list: |
| 415 | return |
| 416 | |
| 417 | try: |
| 418 | col_def = meta_graph_def.collection_def[key] |
| 419 | to_proto = ops.get_to_proto_function(key) |
| 420 | proto_type = ops.get_collection_proto_type(key) |
| 421 | if to_proto: |
| 422 | kind = "bytes_list" |
| 423 | for x in collection_list: |
| 424 | # Additional type check to make sure the returned proto is indeed |
| 425 | # what we expect. |
| 426 | proto = to_proto(x, export_scope=export_scope) |
| 427 | if proto: |
| 428 | assert isinstance(proto, proto_type) |
| 429 | getattr(col_def, kind).value.append(proto.SerializeToString()) |
| 430 | else: |
| 431 | kind = _get_kind_name(collection_list[0]) |
| 432 | if kind == "node_list": |
| 433 | for x in collection_list: |
| 434 | if not export_scope or x.name.startswith(export_scope): |
| 435 | getattr(col_def, kind).value.append( |
| 436 | ops.strip_name_scope(x.name, export_scope)) |
| 437 | elif kind == "bytes_list": |
no test coverage detected