Python multiprocessing + torch tensors: memory problems - memory

Question

I want to generate a dataset (containing numpy arrays) using the python multiprocessing module. I then want to convert the arrays to torch tensors to train a GNN. After each epoch, I want to replace a certain proportion of this dataset with new data. When running the following minimal example of generating and updating the dataset, after each iteration, there are sharp peaks in the memory usage (please see the profile below), until the script eventually crashes due to an oom-kill event.
What are those spikes in the memory usage?
import numpy as np
from multiprocessing import Pool, cpu_count
from torch_geometric.data import Data
import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def generate_buffer():
repeated_args = [arguments] * buffer_size
# create batches in parallel:
with Pool(processes = cpu_count()) as pool:
buffer = pool.starmap(generate_batch, repeated_args)
# flatten the buffer:
buffer = [item for sublist in buffer for item in sublist]
# convert list of numpy arrays to torch Data object containing torch GPU tensors
for i in range(len(buffer)):
X = torch.tensor(buffer[i][0], dtype = torch.float, device = device)
edge_index = torch.tensor(buffer[i][1], dtype = torch.long, device = device)
edge_attr = torch.tensor(buffer[i][2], dtype = torch.float, device = device)
y = torch.tensor(buffer[i][3], dtype = torch.float, device = device)
buffer[i] = Data(x=X, edge_index=edge_index, edge_attr=edge_attr, y=y)
return buffer
def update_buffer(buffer):
# delete the first entries of the buffer:
del buffer[: (replacements_per_iteration * batch_size * len(error_rate))]
# create a list of repeated arguments for all processes:
repeated_args = [arguments] * replacements_per_iteration
# create batches in parallel:
with Pool(processes = cpu_count()) as pool:
new_data = pool.starmap(generate_batch, repeated_args)
# flatten the data:
new_data = [item for sublist in new_data for item in sublist]
# convert list of numpy arrays to torch Data object containing torch GPU tensors
for i in range(len(new_data)):
X = torch.tensor(new_data[i][0], dtype = torch.float, device = device)
edge_index = torch.tensor(new_data[i][1], dtype = torch.long, device = device)
edge_attr = torch.tensor(new_data[i][2], dtype = torch.float, device = device)
y = torch.tensor(new_data[i][3], dtype = torch.float, device = device)
new_data[i] = Data(x=X, edge_index=edge_index, edge_attr=edge_attr, y=y)
# append to buffer:
buffer.extend(new_data)
del new_data
return buffer
def generate_batch(arguments):
batch = []
# need to create a different seed in every thread:
np.random.seed()
for _ in range(batch_size):
graph = generate_sample(arguments)
batch.append(graph)
return batch
def generate_sample(arguments):
# Generating a sample (syndrome measurement of a rotated surface
# code cycle (a quantum error correction scheme)) and mapping
# it to a graph representation using basic numpy operations
return [X, edge_index, edge_attr, y]
if __name__== '__main__':
buffer = generate_buffer()
for i in range(num_iterations):
# update the buffer:
buffer = update_buffer(buffer)
Not converting the data to torch tensors or not using the multiprocessing module (looping over buffer_size) somewhat flattens the spikes. However, I would like to use the speedup of multiprocessing and torch tensors for my application.
No torch tensors
Loop instead of multiprocessing
Just using one process (Pool(processes = 1))
Thanks for your help!