[Bug] Fixes issues in feature pickling and transmission for subgraphs and blocks #2139
Conversation
| def collate(self, items): | ||
| input_nodes, output_nodes, blocks = super().collate(items) | ||
| _pop_blocks_storage(blocks, self.g) | ||
| return input_nodes, output_nodes, blocks |
There was a problem hiding this comment.
Why not directly having _pop_blocks_storage in NodeCollator?
There was a problem hiding this comment.
What if you write your own iterator that gets the blocks with NodeCollator (like in Tensorflow)? The features of the blocks will be removed in this case.
There was a problem hiding this comment.
Or we can put them in the node collator and have an argument pop_storage which defaults to False.
| _pop_subgraph_storage(pair_graph, self.g) | ||
| _pop_subgraph_storage(neg_pair_graph, self.g) | ||
| _pop_blocks_storage(blocks, self.g_sampling) | ||
| return input_nodes, pair_graph, neg_pair_graph, blocks |
There was a problem hiding this comment.
same here. I feel the logic can go to EdgeCollator directly.
| for key, col in subframe._columns.items(): | ||
| if col.storage is None: | ||
| assert key in frame._columns | ||
| col.storage = frame._columns[key].storage |
There was a problem hiding this comment.
This check is only valid when a column's storage can never be None. Maybe add one sentence in Column's docstring to emphasize that?
There was a problem hiding this comment.
I feel the assertion at L63 is redundant, gives no extra information and incurs an extra key lookup. I'd suggest remove it.
There was a problem hiding this comment.
I left one paragraph in the docstring of Column. Does that suffice?
… and blocks (dmlc#2139) * bug fixes * remove __deepcopy__; not sure what the behavior should be * lint * skip gpu test * fix * fix dist dataloader * add comment * remove assert
Description
Fixes #2135 and #2137 by allowing the
storageofColumnto beNone.For subprocesses in PyTorch node and edge DataLoaders, we first check if the feature comes from the original graph by checking whether the column storage of the subgraph/block is identical to that of the original graph. If so, we set the storage to be
None.After the subprocesses transmit the subgraphs/blocks, we then check in the main process if the column storage is None. If so, we restore the storage to be that with the same name in the original graph.