Pickling an optimizer w/ callable attributes

username: GeoffNN
cooked: <p>Hi,</p>
<p>I’m implementing a version of a proximal gradient descent algorithm. For this, I wrote an optimizer which takes a <code>prox</code> callable as argument (e.g. a method which performs soft thresholding, clipping, or some projection/prox) ; this prox is saved as an attribute in the <code>defaults</code> dict.</p>
<p>When I try to save the optimizer’s state dict, I get a pickling error: <code>Attribute error: can't pickle local object 'prox' </code>.</p>
<p>Is there any way to get around this?</p>
<p>Code to reproduce this error: <a href="https://colab.research.google.com/drive/1852mLW9XimED_-mpsz37l-FeWuC6mpqS?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colaboratory</a></p>

username: pascal_notsawo
cooked: <p>Thanks for leaving the link to your notebook</p>
<p>I already had this error when defining one of the methods of my pytorch module as a lambda function, like <code>self.proj = lambda x : x</code>.<br>
When I changed it to <code>self.proj = nn.Identity()</code> I didn’t get any error.<br>
So the error could come from the fact that prox contains lambda functions.</p>
<p>You can change <code>lambda x, s=None: x</code> for example to an <code>nn.Identity()</code>.</p>
<p>Try to remove the lambda functions, and use pytorch module.</p>
<p>Note : I get this when executing your code <code>AttributeError: Can't pickle local object 'PGD.__init__.&lt;locals&gt;.&lt;lambda&gt;'</code></p>

username: GeoffNN
cooked: <p>Ah, thanks <a class="mention" href="/u/pascal_notsawo">@pascal_notsawo</a> – I hadn’t thought of defining the <code>prox</code> operations as <code>nn.Module</code>s. That’s helpful, although will take some time to fix (I have a bunch of constraints defined <a href="https://github.com/openopt/chop/blob/master/chop/constraints.py" class="inline-onebox" rel="noopener nofollow ugc">chop/constraints.py at master · openopt/chop · GitHub</a>).</p>
<p>If another way is possible, I’m all ears.</p>

username: pascal_notsawo
cooked: <p>Here is the modification I made, and it works.<br>
Since prox elements are only seen as functions in the end, there is no need to put them in a ModuleList. But you can still do it.</p>
<pre><code class="lang-auto">class f(torch.nn.Module) :
    def __init__(self, prox_el) :
        super().__init__()
        self.prox_el = prox_el

    def forward(self, x, s=None) :
        return self.prox_el(x.unsqueeze(0)).squeeze()

class PGD(torch.optim.Optimizer):
    """Proximal Gradient Descent
    Args:
      params: [torch.Parameter]
        List of parameters to optimize over
      prox: [callable or None]
        List of prox operators, one per parameter.
      lr: float
        Learning rate
      momentum: float in [0, 1]
      normalization: str
        Type of gradient normalization to be used.
        Possible values are 'none', 'L2', 'Linf', 'sign'.
    """
    name = 'PGD'
    POSSIBLE_NORMALIZATIONS = {'none', 'L2', 'Linf', 'sign'}

    def __init__(self, params, prox=None, lr=.1, momentum=.9, normalization='none'):
        if prox is None:
            prox = [None] * len(list(params))

        self.prox = []
        for prox_el in prox:
            if prox_el is not None:
                self.prox.append(f(prox_el))
            else:
                self.prox.append(torch.nn.Identity())
        #...
        # ...
</code></pre>

username: GeoffNN
cooked: <p>This is awesome; so just subclassing <code>nn.Module</code> and putting the content of the lambda function in the forward.</p>
<p>Thanks! <img src="https://discuss.pytorch.org/images/emoji/apple/partying_face.png?v=9" title=":partying_face:" class="emoji" alt=":partying_face:"></p>

username: pascal_notsawo
cooked: <p>Yes, you are welcome</p>
<p>But my solution is only an avoidance technique, because I think we can simply (and it’s better) define them as simple functions when there are no (learnable) parameters. Except that serialization fails in this case, if we define the function in the ordinary way.</p>
<p>For example, if I define our function above as below the serialization will fail by returning an error like : <code>AttributeError: Can't pickle local object 'f.&lt;locals&gt;.g'</code></p>
<pre><code class="lang-auto">def f(prox_el) :
  def g(x, s=None) :
    return prox_el(x.unsqueeze(0)).squeeze()
  return g
</code></pre>
<p>On the other hand, if in my module I define an attribute like <code>self.loss = F.mse_loss</code>, which here is only a function, the serialization will succeed.</p>
<p>So a solution would be to see how pytorch defines its own functions.</p>