PyTorch: Can't pickle local object 'TorchGraph.create_forward_hook.<locals>.after_forward_hook'

andrebastosdias: <p>I’m working on a multiprocessing reinforcement learning project and have implemented a setup similar to the following example::</p>
<pre><code>if __name__ == '__main__':
    num_processes = 4
    model = MyModel()
    # NOTE: this is required for the ``fork`` method to work
    model.share_memory()
    processes = []
    for rank in range(num_processes):
        p = mp.Process(target=train, args=(model,))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()
</code></pre>
<p>Occasionally, one of the pipes breaks (I suspect this is a server issue, as my workers are surrounded by try clauses but do not throw any errors), and I need to spawn a new process to replace the failed one. However, when I attempt to use p.start() again, I encounter the following error:</p>
<pre><code>AttributeError: Can't pickle local object 'TorchGraph.create_forward_hook.&lt;locals&gt;.after_forward_hook'
</code></pre>
<p>I believe this is happening because the model has already performed forward passes, which prevents it from being pickled. Is there an easy fix for this issue without having to modify the external reinforcement learning algorithm?</p>
<p>Edit: The RL algorithm that I’m using is PPO, I’m overriding the one from <a href="https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/ppo/ppo.py" rel="noopener nofollow ugc">Stable Baselines3</a>.<br>
Edit2: I’m creating the process with</p>
<pre><code>ctx = mp.get_context("forkserver")
ctx.Process(target=train, args=(model,), daemon=True)
</code></pre>

soulitzer: <p>Not sure if either of these are doable or easy, but - Use fork instead of spawn? Make that hook a global function?</p>

andrebastosdias: <ol>
<li>I’m using forkserver (edited main post).</li>
<li>The hook is an internal function of PyTorch, I don’t know how to do that or if it is even possible.</li>
</ol>

soulitzer: <aside class="quote no-group" data-username="andrebastosdias" data-post="1" data-topic="204995">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/andrebastosdias/48/63308_2.png" class="avatar"> andrebastosdias:</div>
<blockquote>
<p><code>TorchGraph</code></p>
</blockquote>
</aside>
<p>I see. Though I think create_forward_hook and after_forward_hook may be internal functions of TorchGraph rather than PyTorch (I wonder if its possible to make the modification there).</p>

andrebastosdias: <p>Oh wait, TorchGraph is from <a href="https://github.com/wandb/wandb/blob/6aa4013daa03ba109815811239e27fae4aa1cfce/wandb/wandb_torch.py#L301" rel="noopener nofollow ugc">Weights and Biases</a>. I just found <a href="https://github.com/wandb/wandb/issues/2588" rel="noopener nofollow ugc">this issue</a> on their github.<br>
I think I can solve the problem by removing this code</p>
<pre><code>        wandb.watch(
            models=self.policy,
            criterion=None,
            log="all",
            log_freq=log_interval,
            log_graph=True
        )
</code></pre>
<p>Thank you <a class="mention" href="/u/soulitzer">@soulitzer</a> !</p>