How to write training loop for MaskRCNN Distributed RPC

matrix: <p>I am trying to train MaskRCNN (class from Torchvision) on 3 machines (master, worker-1 and worker-2) (also multiple GPUs per machine) using distributed RPC &amp; with model parallelism. So far I was able to divide the whole architecture onto 2 machines (manually). So far I have used below resources extensively to do this,</p>
<p>Finetuning MaskRCNN in general: <a href="https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">TorchVision Object Detection Finetuning Tutorial — PyTorch Tutorials 1.7.1 documentation</a><br>
Distributed RPC: <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Distributed Pipeline Parallelism Using RPC — PyTorch Tutorials 1.7.1 documentation</a><br>
Torchvision github</p>
<p>Now only thing left is to write the training loop, where I am stuck. I have observed the training loop used in <a href="https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html" rel="noopener nofollow ugc">this</a>. And I need to convert this code to distributed training loop. How should I edit <a href="https://github.com/pytorch/vision/tree/master/references/detection/engine.py" rel="noopener nofollow ugc">https://github.com/pytorch/vision/tree/master/references/detection/engine.py</a> (Which contains training loop for training mask-rcnn as linked previously) to achieve this? Also training loop for distributed rpc is also given in <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html" rel="noopener nofollow ugc">this example</a> (As mentioned above) for classification task. How to combine these 2 to train mask rcnn in distributed way (with model parallelism)?</p>
<p>EDIT :</p>
<p>Question-2:</p>
<p>Second shard of the model which resides on the second machine, has 2 modules. Each of them is allocated 1 GPU on that machine (meaning module 1 on GPU-1, module 2 on GPU-2). And I have other 2 GPUs on the same machine, which are ideal. Dividing these 2 modules on 4 GPUs equally mean overriding their forward methods, which are very complicated (Talking about RPN and RoI Heads module). Can I do something like this, put module-1 on GPU-1 &amp; GPU-2, module-2 on GPU-3 &amp; GPU-4, now when input batch comes split that equally into 2 parts, process them like this, part1 → module-1 (GPU:1) → module-2 (GPU:3) &amp; part2 → module-1 (GPU:2) → module-2 (GPU:4) (of course, these 2 workflows will run concurrently, at the end their losses will be averaged and avg. loss will be returned to master node). I have already referred to following articles,</p>
<p><a href="https://pytorch.org/tutorials/advanced/rpc_ddp_tutorial.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/advanced/rpc_ddp_tutorial.html</a></p>
<p><a href="https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html</a></p>
<p>Is it possible to perform above explained case in PyTorch?</p>

wayi: <p>To be clear, since you already know how to split the model into two model shards, is your question how to convert the following lines in <a href="https://github.com/pytorch/vision/blob/master/references/detection/engine.py:" rel="noopener nofollow ugc">https://github.com/pytorch/vision/blob/master/references/detection/engine.py:</a></p>
<pre><code class="lang-auto">        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
</code></pre>
<p>into the code in <a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html#step-3-define-the-training-loop" rel="noopener nofollow ugc">RPC Framework Tutorial</a>?</p>
<pre><code class="lang-auto">        with dist_autograd.context() as context_id:
            outputs = model(inputs)
            dist_autograd.backward(context_id, [loss_fn(outputs, labels)])
            opt.step(context_id)
</code></pre>
<p>cc: <a class="mention" href="/u/mrshenli">@mrshenli</a></p>

matrix: <p>I know about the conversion you mentioned in your answer. But what should be done to these lines in engine.py. Should I keep it as it is?</p>
<pre><code>    # reduce losses over all GPUs for logging purposes
    loss_dict_reduced = utils.reduce_dict(loss_dict)
    losses_reduced = sum(loss for loss in loss_dict_reduced.values())
</code></pre>
<p>I have added another question please take look at that. Thanks for answering my question<strong>s</strong>.</p>

mrshenli: <blockquote>
<p>I know about the conversion you mentioned in your answer. But what should be done to these lines in engine.py. Should I keep it as it is?</p>
</blockquote>
<p>Hey <a class="mention" href="/u/matrix">@matrix</a>, <code>reduce_dict</code> internally uses collective communication. If you need that, you will need to setup the process group with <code>init_process_group</code> properly in the way that only the processes that hold <code>loss_dict</code> form a gang.</p>
<p>However, since you already use RPC, do you still need <code>reduce_dict</code>? Will it work if you send the loss to the same process, reduce those loss their locally (without using collective communication), and then launch distributed backward from there?</p>