Conflicts between torch-1.10.2+cu113 and python multiprocessing

username: SixerWang
cooked: <pre data-code-wrap="python"><code class="lang-python">import os, re, sys, time
import numpy as np
from concurrent.futures import ProcessPoolExecutor
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(['0', '1', '2', '3'])

class Model(nn.Module):

    def __init__(self, params=None):
        super().__init__()
        input_shape = params.get('input_size')
        output_shape = params.get('output_size')
        self.gru = nn.GRU(input_shape, output_shape, batch_first=True) # Note

    def forward(self, x):
        y = self.gru(x)
        return y

# sub-process job
def single_job(cuda_id, params):
    device = torch.device(cuda_id)
    model = Model(params)
    model.to(device)
    for _ in range(1000):
        X = torch.rand(4096, 16, 64).float()
        X = X.to(device)
        y_pred = model(X) # Note
        time.sleep(1)

# main process
with ProcessPoolExecutor(max_workers=2) as executor:
    for i in [1, 2]:
        cuda_id = f'cuda:{i}'
        params = {'input_size': 64, 'output_size': 1}
        sub = executor.submit(single_job, cuda_id, params)
</code></pre>
<p>However, when it runs, the GPU occupation could be weird as shown below.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f29d63e887f623417ea8740c1de9b4f360b536a.png" data-download-href="https://discuss.pytorch.org/uploads/default/1f29d63e887f623417ea8740c1de9b4f360b536a" title="256af90e439b73b3a0c1569eb0b6c8e"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1f29d63e887f623417ea8740c1de9b4f360b536a_2_568x500.png" alt="256af90e439b73b3a0c1569eb0b6c8e" data-base62-sha1="4rGr4IrN3RGRB2MBUIGJqC8RzIu" width="568" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1f29d63e887f623417ea8740c1de9b4f360b536a_2_568x500.png, https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f29d63e887f623417ea8740c1de9b4f360b536a.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/f/1f29d63e887f623417ea8740c1de9b4f360b536a.png 2x" data-dominant-color="0E0E0D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">256af90e439b73b3a0c1569eb0b6c8e</span><span class="informations">850×748 180 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>At the begining, the gpu occupation shown in nvidia-smi page is normal, with gpu 1 and 2 are occupied separately.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7d904d7e0373036b7215d2890de0e4fe4cf1e547.png" data-download-href="https://discuss.pytorch.org/uploads/default/7d904d7e0373036b7215d2890de0e4fe4cf1e547" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7d904d7e0373036b7215d2890de0e4fe4cf1e547.png" alt="image" data-base62-sha1="hUMQKWWQkaeEkjv9Y3OdOzK6mj5" width="553" height="500" data-dominant-color="0E0E0D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">721×651 34.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>But after a few seconds, another two process is generated and occupy gpu 0 unexpectedly.<br>
This problem is only caused by GRU layer (including other RNN layer) in model propagation as shown in the source code above.</p>
<p>The code runs correctly with torch-1.7 and corresponding cuda version.</p>