How does one pickle arbitrary pytorch models that use lambda functions?

username: pinocchio
cooked: <p>I currently have a neural network module:</p>
<pre><code class="lang-auto">import torch.nn as nn

class NN(nn.Module):
    def __init__(self,args,lambda_f,nn1, loss, opt):
        super().__init__()
        self.args = args
        self.lambda_f = lambda_f
        self.nn1 = nn1
        self.loss = loss
        self.opt = opt
        # more nn.Params stuff etc...

    def forward(self, x):
        #some code using fields
        return out
</code></pre>
<p>I am trying to checkpoint it but because pytorch saves using <code>state_dict</code>s it means I can’t save the lambda functions I was actually using if I checkpoint with the pytorch <code>torch.save</code> etc. I literally want to save everything without issue and re-load to train on GPUs later. I currently am using this:</p>
<pre><code class="lang-auto">def save_ckpt(path_to_ckpt):
    from pathlib import Path
    import dill as pickle
    ## Make dir. Throw no exceptions if it already exists
    path_to_ckpt.mkdir(parents=True, exist_ok=True)
    ckpt_path_plus_path = path_to_ckpt / Path('db')

    ## Pickle args
    db['crazy_mdl'] = crazy_mdl
    with open(ckpt_path_plus_path , 'ab') as db_file:
        pickle.dump(db, db_file)
</code></pre>
<p>currently it throws no errors when I chekpoint it and it saved it.</p>
<p>I am worried that when I train it there might be a subtle bug even if no exceptions/errors are trained or something unexpected might happen (e.g. weird saving on disks in the clusters etc who knows).</p>
<p>Is this safe to do with pytorch classes/nn models? Especially if we want to resume training with GPUs?</p>
<hr>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/61510810/how-does-one-pickle-arbitrary-pytorch-models-that-use-lambda-functions" target="_blank" rel="nofollow noopener">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/3167448/pinocchio" target="_blank" rel="nofollow noopener">
    <img alt="Pinocchio" src="https://www.gravatar.com/avatar/ba92064b084638cb903cbd271120be64?s=128&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="128" height="128">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/61510810/how-does-one-pickle-arbitrary-pytorch-models-that-use-lambda-functions" target="_blank" rel="nofollow noopener">How does one pickle arbitrary pytorch models that use lambda functions?</a>
</h4>

<div class="tags">
  <strong>python, pytorch, pickle</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/3167448/pinocchio" target="_blank" rel="nofollow noopener">
    Pinocchio
  </a>
  on <a href="https://stackoverflow.com/questions/61510810/how-does-one-pickle-arbitrary-pytorch-models-that-use-lambda-functions" target="_blank" rel="nofollow noopener">08:11PM - 29 Apr 20 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Cross posted:</p>
<ul>
<li><a href="https://stackoverflow.com/questions/61510810/how-does-one-pickle-arbitrary-pytorch-models-that-use-lambda-functions" rel="nofollow noopener">https://stackoverflow.com/questions/61510810/how-does-one-pickle-arbitrary-pytorch-models-that-use-lambda-functions</a></li>
<li><a href="https://discuss.pytorch.org/t/how-does-one-pickle-arbitrary-pytorch-models-that-use-lambda-functions/79026" class="inline-onebox">How does one pickle arbitrary pytorch models that use lambda functions?</a></li>
<li>
<a href="https://www.reddit.com/r/pytorch/comments/gagpjg/how_does_one_pickle_arbitrary_pytorch_models_that/" rel="nofollow noopener">https://www.reddit.com/r/pytorch/comments/gagpjg/how_does_one_pickle_arbitrary_pytorch_models_that/</a>?</li>
<li><a href="https://www.quora.com/unanswered/How-does-one-pickle-arbitrary-PyTorch-models-that-use-lambda-functions" rel="nofollow noopener">https://www.quora.com/unanswered/How-does-one-pickle-arbitrary-PyTorch-models-that-use-lambda-functions</a></li>
</ul>
<hr>
<p>related: <a href="https://discuss.pytorch.org/t/advantages-disadvantages-of-using-pickle-module-to-save-models-vs-torch-save/79016" class="inline-onebox">Advantages &amp; Disadvantages of using pickle module to save models vs torch.save</a></p>

username: GeoffNN
cooked: <p>In case this is still an issue, we have a (partly satisfactory) solution <a href="https://discuss.pytorch.org/t/pickling-an-optimizer-w-callable-attributes/120527/5">here</a>.</p>