Detach rsamples from distributions parameters

iacolippo: <p>I want to compute the gradient for the parameters of a distributions in two steps, such that it is possible to decouple the code defining the distribution from the for loop training. The following works but I have to retain the whole graph (or manually delete references to <code>cost</code> and <code>samples</code> with <code>del</code>, which may become unfeasible when the computational graph becomes more complex).</p>
<pre><code class="lang-auto">log_std = nn.Parameter(torch.Tensor([1]))
std = torch.exp(log_std)
mean = nn.Parameter(torch.Tensor([1]))
dist = dd.Normal(loc=mean, scale=std)
optim = torch.optim.SGD([log_std, mean], lr=0.01)
target = dd.Normal(5,5)

for i in range(50):
    optim.zero_grad()
    samples = dist.rsample((1000,))
    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum()
    cost.backward(retain_graph=True)
    optim.step()
    print(i, log_std, mean, cost)
</code></pre>
<p>I would like to do something like this instead</p>
<pre><code class="lang-auto">log_std = nn.Parameter(torch.Tensor([1]))
std = torch.exp(log_std)
mean = nn.Parameter(torch.Tensor([1]))
dist = dd.Normal(loc=mean, scale=std)

optim = torch.optim.SGD([log_std, mean], lr=0.01)
target = dd.Normal(5,5)

for i in range(50):
    optim.zero_grad()
    samples = dist.rsample((1000,))
    detached_samples = samples.detach()
    detached_samples.retain_grad()
    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum()
    cost.backward() # compute gradients up to detached samples
    samples.backward(detached_samples.grad(), retain_graph=True)  # get gradients of distribution parameters log_std, mean and retain only this part of the graph
    optim.step()  # apply update
    print(i, log_std, mean, cost)
</code></pre>
<p>Even if I detach samples, when I run <code>cost.backward()</code> I get gradients for <code>log_std</code> and <code>mean</code>. I don’t know if that’s intended and what I’m doing is smart or dumb or even possible, but I would like someone to shed some light on this <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

tom: <p>Well, if you use samples rather than detached_samples in your cost…</p>

iacolippo: <p>oh my god, that was dumb <img src="https://discuss.pytorch.org/images/emoji/apple/laughing.png?v=9" title=":laughing:" class="emoji" alt=":laughing:"></p>
<p>using</p>
<pre><code class="lang-auto">cost = -(target.log_prob(detached_samples) - dist.log_prob(detached_samples)).sum()
</code></pre>
<p>doesn’t solve the problem though <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"> I think it’s impossible to do this because there is a direct dependence between log_prob and the parameters of <code>Normal</code>, detaching samples doesn’t actually forbid the gradients to flow up to <code>log_std</code> and <code>mean</code>.</p>
<p>Thanks for taking a look!</p>

tom: <p>So one other thing is that you likely want <code>detached_samples.requires_grad_()</code>.<br>
So you want gradients w.r.t. detached_samples without affecting the params? Then I would recommend using <code>torch.autograd.grad</code> instead of <code>.backward</code> to selectively calculate gradients without side-effects.</p>

Michael_Hutchinson: <p>So this question came from a <a href="https://stackoverflow.com/questions/56955932/" rel="nofollow noopener">stack overflow post</a> of mine; The ultimate goal is to train the parameters of the distribution, without having to worry about<br>
a) the type of distribution used, or<br>
b) how the parameters of the distribution are generated (in this case, the std of the normal distribution is defined in terms of its log to ensure non-negative results)</p>
<p>In short, here this means I do not want to have to do std = e^(log_std) at any point other than the first time I define the distribution.</p>
<p>This is something that will work in Tensorflow, as the graphs are static there. It would seem the easiest way would be to retain some small parts of the graph in PyTorch, but this seems to not be possible. The above was a way of trying to achieve that. <a class="mention" href="/u/tom">@tom</a> If you have any suggestions on this I would greatly appreciate it!</p>

tom: <aside class="quote no-group" data-username="Michael_Hutchinson" data-post="5" data-topic="50251">
<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/michael_hutchinson/48/13913_2.png" class="avatar"> Michael_Hutchinson:</div>
<blockquote>
<p>This is something that will work in Tensorflow, as the graphs are static there.</p>
</blockquote>
</aside>
<p>Well yeah, but TF will <em>do</em> the computation over and over again.<br>
So what you would do in PyTorch is to wrap parameter, transformation of the parameter and distribution in a <code>nn.Module</code> and use that.<br>
A long while ago, I thought that blurring the boundary between <code>nn.Modules</code> and <code>nn.Parameters</code> <a href="https://github.com/pytorch/pytorch/issues/7313" rel="noopener nofollow ugc">was a good idea</a>, I still think it would be, but obviously it doesn’t seem to be immediately pending.</p>
<p>Best regards</p>
<p>Thomas</p>

Michael_Hutchinson: <p>Thank you <a class="mention" href="/u/tom">@tom</a>, do you happen to have a quick hints on implementing this? The main issue I see is if I pass a module in the place of a parameter, it will not be accesses in the same way (i.e. self.parameter vs self.parameter() ).</p>
<p>Thanks,<br>
Michael</p>

tom: <p>No, don’t use the Module for parameters (doesn’t work well, or we would not need the feature discussed in the linked abstraction), but have the nn.Module wrapping your distribution. You can have your module define log_prob and rsample if you want - except for calling some hooks that you might not need between model() and model.forward(), forward isn’t special.<br>
In a way, this is similar to the <a href="https://pytorch.org/tutorials/beginner/nn_tutorial.html#closing-thoughts" rel="nofollow noopener">what is nn.Module</a> tutorial - nn.Modules are containers for the parameters and wrap the calculation we do with them.</p>
<p>Best regards</p>
<p>Thomas</p>

neerajprad: <p>Constructing new distribution instances is cheap, and you can use the <code>constraints</code> library to automatically handle the transformation of the <code>std</code> parameter to <code>(0, inf)</code>, as below.</p>
<p><strong>Code:</strong></p>
<pre><code class="lang-python">import torch
import torch.distributions as dd
from torch.distributions import biject_to

params = {'loc': torch.tensor([1.], requires_grad=True),
          'scale': torch.tensor([1.], requires_grad=True)}

optim = torch.optim.Adam(params.values(), lr=0.001)
target = dd.Normal(5, 5)


def get_dist():
    return dd.Normal(**get_constrained_params())


def get_constrained_params():
    # apply automatic constraints to distribution parameters.
    constraints = target.arg_constraints
    constrained_params = {}
    for k, v in params.items():
        constrained_params[k] = biject_to(constraints[k])(v)
    return constrained_params


for i in range(10000):
    optim.zero_grad()
    dist = get_dist()
    samples = dist.rsample((1000,))
    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum()
    cost.backward()
    optim.step()
    print(i, get_constrained_params())

</code></pre>
<p>Does this satisfy what you had in mind?</p>

tom: <p>That’s awesome, <a class="mention" href="/u/neerajprad">@neerajprad</a>! I’d probably stick the things into a nn.Module to hold the parameters. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Best regards</p>
<p>Thomas</p>

neerajprad: <p>Based on <a class="mention" href="/u/tom">@tom</a>’s suggestion, here is a modified attempt that uses <code>nn.Module</code> <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"></p>
<pre><code class="lang-python">import torch
import torch.distributions as dd
import torch.nn as nn
from torch.distributions import biject_to


class DistributionModule(nn.Module):
    def __init__(self, dist, arg_constraints, **params):
        super(DistributionModule, self).__init__()
        self.dist_class = dist
        self.arg_constraints = arg_constraints
        for k, v in params.items():
            self.register_parameter(k, nn.Parameter(v))

    def get_dist(self):
        constrained_params = dict(self.get_constrained_params())
        return self.dist_class(**constrained_params)

    def log_prob(self, sample):
        return self.get_dist().log_prob(sample)

    def rsample(self, sample_shape=()):
        return self.get_dist().rsample(sample_shape)

    def forward(self, value):
        return self.log_prob(value)

    def get_constrained_params(self):
        for name, param in self.named_parameters():
            yield name, biject_to(self.arg_constraints[name])(param)


params = {'loc': torch.tensor([1.]),
          'scale': torch.tensor([1.])}
target = dd.Normal(5, 5)
d = DistributionModule(dd.Normal, target.arg_constraints, **params)
optim = torch.optim.Adam(d.parameters(), lr=0.001)


for i in range(10000):
    optim.zero_grad()
    samples = d.rsample((1000,))
    cost = -(target.log_prob(samples) - d.log_prob(samples)).sum()
    cost.backward()
    optim.step()
    print(i, dict(d.get_constrained_params()))
</code></pre>
<p>Not tested very well, but I think you should be able to use a generic wrapper like this for most distribution classes! The forward method just delegates to <code>.log_prob</code> (though it is not being used here).</p>