'How to get gpu data from torch.multiprocessing.Queue?
Here is my code:
import torch
import torch.multiprocessing as mp
mp.set_start_method('spawn', force=True)
def job(device, q, event):
x = torch.ByteTensor([1,9,5]).to(device)
x.share_memory_()
print("in job:", x)
q.put(x)
event.wait()
def main():
device = torch.device("cuda" if torch.cuda.is_available else "cpu")
num_processes = 4
processes = []
q = mp.SimpleQueue()
event = mp.Event()
for rank in range(num_processes):
p = mp.Process(target=job, args=(device, q, event))
p.start()
processes.append(p)
for i in range(num_processes):
y = q.get()
print("in main:", y)
del y
event.set()
for p in processes:
p.join()
if __name__ =="__main__":
main()
Got the outputs:
in job: tensor([1, 9, 5], device='cuda:0', dtype=torch.uint8)
in job: tensor([1, 9, 5], device='cuda:0', dtype=torch.uint8)
in job: tensor([1, 9, 5], device='cuda:0', dtype=torch.uint8)
in job: tensor([1, 9, 5], device='cuda:0', dtype=torch.uint8)
in main: tensor([0, 0, 0], device='cuda:0', dtype=torch.uint8)
in main: tensor([0, 0, 0], device='cuda:0', dtype=torch.uint8)
in main: tensor([0, 0, 0], device='cuda:0', dtype=torch.uint8)
in main: tensor([0, 0, 0], device='cuda:0', dtype=torch.uint8)
When tensor is on cpu, good.But when I send it to CUDA, I can't read the tensor content from main process.What triggers gpu data lost through queue? I think it could be some mechanism of multiprocessing memory space but Im not quite familiar with it.
Sources
This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.
Source: Stack Overflow
Solution | Source |
---|