Unfortunately, for quite some time now, I have encountered problems with the module torch.multiprocessing. As a MWE, I am trying to square a PyTorch tensor on CPU, which does not work:
import torch
import numpy as np
import torch.multiprocessing as mp
def square(i, x, queue):
print('In process {}'.format(i))
if isinstance(x, np.ndarray):
queue.put(np.square(x))
else:
queue.put(torch.square(x))
if __name__ == '__main__':
mp.set_start_method('spawn')
print('Number of available CPUs: {}'.format(mp.cpu_count()))
processes = []
queue = mp.Queue()
# Set up some sample data:
x = np.arange(64)
# Start <mp.cpu_count()> processes with the square function as the target and an invidual piece of data to process:
for i in range(mp.cpu_count()):
start_index = mp.cpu_count()*i
proc = mp.Process(target = square, args = (i, x[start_index:start_index+mp.cpu_count()], queue))
proc.start()
processes.append(proc)
# Wait for each process to finish before returning to the main thread:
for proc in processes:
proc.join()
# Terminate each process:
for proc in processes:
proc.terminate()
# Convert the multiprocessing queue into a list:
results = []
while not queue.empty():
results.append(queue.get())
print(results)
This results in the following error message (Python version 3.8.9, PyTorch version 1.9.0):
Traceback (most recent call last):
File "parallelization_sampling.py", line 55, in <module>
results.append(queue.get())
File "/..../lib/python3.8/multiprocessing/queues.py", line 116, in get
return _ForkingPickler.loads(res)
File "/..../lib/python3.8/site-packages/torch/multiprocessing/reductions.py", line 289, in rebuild_storage_fd
fd = df.detach()
File "/..../lib/python3.8/multiprocessing/resource_sharer.py", line 57, in detach
with _resource_sharer.get_connection(self._id) as conn:
File "/.../lib/python3.8/multiprocessing/resource_sharer.py", line 87, in get_connection
c = Client(address, authkey=process.current_process().authkey)
File "/.../lib/python3.8/multiprocessing/connection.py", line 502, in Client
c = SocketClient(address)
File "/.../lib/python3.8/multiprocessing/connection.py", line 630, in SocketClient
s.connect(address)
ConnectionRefusedError: [Errno 111] Connection refused
However, when I import the module multiprocessing instead of torch.multiprocessing and use x = np.arange(64) instead of x = torch.arange(64), the script works fine! I tested the code on two machines, can anybody confirm this weird behavior? And does anybody see what I am doing wrongly? Or is torch.multiprocessing not supposed to be used on CPU?