Hi, I am trying to train several models in parallel using torch 's pool.map(). Since my setup has multiple GPUs, I pass a device also to my training task and the model is trained on that particular device.
The problem I face is RuntimeError: CUDA error: out of memory
after a while.
This happens after several models are trained and I can clearly see using watch nvidia-smi that the GPU memory accumulates over time.
I have posted a minimal example below which also leads to the stated issue.
Is there something needs to be done at the end of run() to clean up the memory?
Pytorch version: torch 1.6.0
import torch
from torch.multiprocessing import Pool
import random
class SimpleModule(torch.nn.Module):
def __init__(self, input_: int, output: int):
super().__init__()
self.linear = torch.nn.Linear(input_, output)
def forward(self, input_: torch.Tensor):
return self.linear(input_)
def run(device: str):
model = SimpleModule(10, 10).to(device)
optimizer = torch.optim.SGD(lr=1e+2,
params=model.parameters())
loss_criterion = torch.nn.MSELoss()
for i in range(int(1e+4)):
optimizer.zero_grad()
inputs = torch.rand(5, 10).to(device)
outputs = torch.rand(5, 10).to(device)
preds = model.forward(inputs)
loss = loss_criterion(preds, outputs)
loss.backward()
optimizer.step()
if __name__ == "__main__":
torch.multiprocessing.set_start_method("spawn", force=True)
n_tasks = 500
available_device = ["cpu"] if not torch.cuda.is_available() else [f"cuda:{i}" for i in
range(torch.cuda.device_count())]
devices = random.choices(available_device, k=n_tasks)
pool = Pool(processes=50)
pool.map(run, devices)