Removes the storage objects in the given graphs' Frames if it is a sub-frame of the given parent graph, so that the storages are not serialized during IPC from PyTorch DataLoader workers.
(item, g)
| 574 | # DGLGraphs will preserve the lazy feature slicing for subgraphs. (2) Other graph storages |
| 575 | # will not have lazy feature slicing; all feature slicing will be eager. |
| 576 | def remove_parent_storage_columns(item, g): |
| 577 | """Removes the storage objects in the given graphs' Frames if it is a sub-frame of the |
| 578 | given parent graph, so that the storages are not serialized during IPC from PyTorch |
| 579 | DataLoader workers. |
| 580 | """ |
| 581 | if not isinstance(item, DGLGraph) or not isinstance(g, DGLGraph): |
| 582 | return item |
| 583 | |
| 584 | for subframe, frame in zip( |
| 585 | itertools.chain(item._node_frames, item._edge_frames), |
| 586 | itertools.chain(g._node_frames, g._edge_frames), |
| 587 | ): |
| 588 | for key in list(subframe.keys()): |
| 589 | subcol = subframe._columns[key] # directly get the column object |
| 590 | if isinstance(subcol, LazyFeature): |
| 591 | continue |
| 592 | col = frame._columns.get(key, None) |
| 593 | if col is None: |
| 594 | continue |
| 595 | if col.storage is subcol.storage: |
| 596 | subcol.storage = None |
| 597 | return item |
| 598 | |
| 599 | |
| 600 | def restore_parent_storage_columns(item, g): |