CUDA error: initialization error when using DataLoader with num_workers>0

Xin-yang-Liu: <p>I defined a dataset with a time-consuming index converting function so I want to set <code>num_workers</code> to a large number to accelerate it. However, the dataset contains part of the data that is already on the GPU (to avoid transfer from system memory to GPU memory) and the Dataloader will raise <code>CUDA error: initialization error</code> when <code>num_workers&gt;0</code></p>
<p>I read some posts said people should avoid multiprocessing when dealing with CUDA tensors.  Are there any other ways to make the dataloader faster?</p>

Xin-yang-Liu: <p>Here is the Dataset I defined. The method <code>__index()</code> is used to convert the index and seems slow.</p>
<pre><code class="lang-python">class myset(Dataset):
    def __init__(self, 
                 u: torch.Tensor, 
                 p: torch.Tensor,
                 length = (2, 3.2, 3.2),
                 num_points:int=216*101*257*257,
                 batch_size:int=50000,
                 device:torch.device=torch.device('cpu')):
        self.N = num_points
        
        self.num_p = u.shape[0]
        self.num_t = u.shape[1]
        self.Nx = u.shape[3]
        self.Ny = u.shape[4]
        self.num_f = self.Nx*self.Ny
        self.num_tf = self.num_t*self.num_f
        
        self.p = p.to(device)

        xmesh = torch.linspace(0, length[1], self.Nx, device=device)
        ymesh = torch.linspace(0, length[2], self.Ny, device=device)
        xmesh, ymesh = torch.meshgrid(xmesh, ymesh, indexing='ij')
        self.xmesh, self.ymesh = xmesh.reshape(-1,1), ymesh.reshape(-1,1)
        self.t = torch.linspace(0, length[0], self.num_t, device=device).reshape(-1,1)

        self.uicin = u[:,0].to(device)

        self.num_initp = u.shape[0]
        self.length = length
        self.device = device
        self.batch_size = batch_size
        
        '''
                 0
            +---------+
            |         |
          3 |         | 1
            |         |
            +---------+
                 2
        '''
        
        u = torch.permute(u,(0,1,3,4,2)) 
        self.u = u.reshape(-1, u.shape[4])
        self.uicout = u[:,0].reshape(self.num_p, self.num_f, 2)#.to(device)
        
    def __index(self, i):
        idp = i//self.num_tf
        idf = i%self.num_f
        idt = (i-idp*self.num_tf)//self.num_f
        return idp, idt, idf


    def __getitem__(self, index):
        idp, idt, idf = self.__index(index)
        uin = self.uicin[idp]
        pin = self.p[idp]
        ulable = self.u[index]

        x = self.xmesh[idf]
        y = self.ymesh[idf]
        t = self.t[idt]

        return (x, y, t, uin, pin), ulable


    def __len__(self):
        return self.N
</code></pre>

ptrblck: <p>You might want to change the “start method” for <code>mp</code> as described <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html#cuda-in-multiprocessing">here</a> as each process is currently trying to reinitialize a CUDA context and fails.</p>

Xin-yang-Liu: <p>Thanks for the suggestion. I tried <code>spawn</code> and <code>forkserver</code> but none of them work.</p>
<ul>
<li>
<p>In the <code>spawn</code> mode, I got the same same error “CUDA error: initialization error”, though I did see several subprocesses running before it exit with the error.</p>
</li>
<li>
<p>In the <code>forkserver</code> mode, I got the error:</p>
</li>
</ul>
<pre><code class="lang-auto">RuntimeError: Attempted to set the storage of a tensor on device "cuda:1" to a storage on different device "cuda:0".  This is no longer allowed; the devices must match.

[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]
</code></pre>
<p>However, I only used <code>cuda:1</code> and <code>cpu</code>. Do you have any suggestions for this? Thanks in advance.</p>