[TOC]

  1. Title: Pytorch Multiprocessing 2023
  2. Review Date: Tue, Jul 18, 2023
  3. url: https://pytorch.org/docs/stable/multiprocessing.html

Best Practices for Pytorch multiprocessing

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
from torch import multiprocessing
from torch.multiprocessing import Pipe

multiprocessing.set_start_method('spawn') # for torch multiprocessing, spawn method is required 

model = Model()
model.share_memory()

p_conn, c_conn = Pipe()

tensor_1 = torch.tensor([1])
tensor_1.share_memory() # need to pass tensor to shared memeory 

p_conn.send(tensor_1)


# the receiver end 

tensor_1_sub = c_conn.receive()
# after using it 
del tensor_1_sub # safe del