Severin Magel
Severin Magel

Reputation: 11

How to investigate memory consumption of pytorch_geometric data

I am working on a framework that uses pytorch_geometric graph data stored in the usual way in data.x and data.edge_index Additionally, the data loading process appends multiple other keys to that data object, such as the path to the database or the model's name, both as strings. Now, I would like to see which of those additional fields in the data has how much memory consumption. The goal is to slim those data representations down to increase the batch size while training.

I know that within pytorch geometric, there is the function get_data_size, but it only displays the total theoretical memory consumption. I am also unsure what "theoretical" means in this case.

I`ve tried to do this to see the difference in memory consumption when deleting a key in data, but for the fields with strings in them, this gave 0, which does not make sense to me.

for key in data.keys():
    start = get_data_size(data)
    print(start)
    del data[key]
    end = get_data_size(data)
    print(f"Safed: {start-end} by deleteing {key}")

Upvotes: 1

Views: 28

Answers (0)

Related Questions