How sharing_memory actually worked in pytorch

username: Forceless
cooked: <p>Sorry to bother you, I recently work with torch.multiprocessing and have some problems，hope you can help me😶‍🌫️</p>
<ol>
<li>If I shared a model in cuda, it raises<br>
<code>RuntimeError: Attempted to send CUDA tensor received from another process; this is not currently supported. Consider cloning before sending.</code></li>
<li><code>torch.multiprocessing.manager.queue.get</code> taken a long time to finish. If the queue just passed a file descriptor, I don’t think it should take 1/3 of the total time, is there any faster way?<br>
Here’s my script</li>
</ol>
<pre data-code-wrap="python"><code class="lang-python">import torch

import torch.multiprocessing as mp

from copy import deepcopy

from functools import partial

from time import *

from torchvision import models

import numpy as np

from tqdm import tqdm

def parallel_produce(

    queue: mp.Queue,

    model_method,

    i

) -&gt; None:

    pure_model: torch.nn.Module = model_method()

    # if you delete this line, model can be passed
    pure_model.to('cuda')

    pure_model.share_memory()

    while True:

        corrupt_model = deepcopy(pure_model)

        dic = corrupt_model.state_dict()

        dic[list(dic.keys())[0]]*=2

        corrupt_model.share_memory()

        queue.put(corrupt_model)

def parallel(

    valid,

    iteration: int = 1000,

    process_size: int=2,

    buffer_size: int=2

):

    pool = mp.Pool(process_size)

    manager = mp.Manager()

    queue = manager.Queue(buffer_size)

    SeedSequence = np.random.SeedSequence()

    model_method = partial(models.squeezenet1_1,True)

    async_result = pool.map_async(

        partial(

            parallel_produce,

            queue,

            model_method,

        ),

        SeedSequence.spawn(process_size),

    )

    time = 0

    for iter_times in tqdm(range(iteration)):

        start = monotonic_ns()

        # this takes a long time

        corrupt_model: torch.nn.Module = queue.get()

        time += monotonic_ns() - start

        corrupt_model.to("cuda")

        corrupt_result = corrupt_model(valid)

        del corrupt_model

    pool.terminate()

    print(time / 1e9)

if __name__ == "__main__":

    valid = torch.randn(1,3,224,224).to('cuda')

    parallel(valid)

</code></pre>
<p><span class="hashtag-raw">#total</span> time of queue.get taken</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/0/90b45f4a2a0a06e16c8b8e87654bf7d28b77e4ab.png" data-download-href="https://discuss.pytorch.org/uploads/default/90b45f4a2a0a06e16c8b8e87654bf7d28b77e4ab" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/90b45f4a2a0a06e16c8b8e87654bf7d28b77e4ab_2_690x25.png" alt="image" data-base62-sha1="kE7cuorBUVWwOLcW2SMh2jqJsL1" width="690" height="25" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/90b45f4a2a0a06e16c8b8e87654bf7d28b77e4ab_2_690x25.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/90b45f4a2a0a06e16c8b8e87654bf7d28b77e4ab_2_1035x37.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/90b45f4a2a0a06e16c8b8e87654bf7d28b77e4ab_2_1380x50.png 2x" data-dominant-color="585858"><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">1740×65 8.01 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>