Theano library function equivalent in PyTorch

username: ahmadreza9
cooked: <p>Hello,<br>
If you see (<a href="http://deeplearning.net/software/theano/library/compile/function.html" rel="nofollow noopener">http://deeplearning.net/software/theano/library/compile/function.html</a>) Theano can create functions. I absolutely sure that PyTorch also can have something like that but how we can equivalent the Update parameter of Theano function in PyTorch? (I saw this <a href="https://discuss.pytorch.org/t/porting-code-from-thano-lasagne-to-pytorch/36388" class="inline-onebox">Porting code from Thano/Lasagne to PyTorch</a> but they didn’t talk about the update)</p>
<p>Let’ s have an example:</p>
<pre><code class="lang-auto">def rmsprop_updates(grads, params, stepsize, rho=0.9, epsilon=1e-9):

    updates = []

    for param, grad in zip(params, grads):
        accum = theano.shared(np.zeros(param.get_value(borrow=True).shape, dtype=param.dtype))
        accum_new = rho * accum + (1 - rho) * grad ** 2
        updates.append((accum, accum_new))
        updates.append((param, param + (stepsize * grad / T.sqrt(accum_new + epsilon))))
        # lasagne has '-' after param
    return updates

updates = rmsprop_updates( grads, params, self.lr_rate, self.rms_rho, self.rms_eps)

N = states.shape[0]
loss = T.log(prob_act[T.arange(N), actions]).dot(values) / N

self._train_fn = theano.function([states, actions, values], loss, updates=updates)

</code></pre>

username: iffiX
cooked: <p>Instead of compiling a function from a string, the programming paradigm used by pytorch is using modules as building blocks to create a function, there is only one important thing in pytorch: <strong>tensor</strong>, tensor is both the container for data, the interface for low-level device handling, and the gradient flow component. Therefore, instead of defining a graph, you construct a graph <strong>dynamically</strong> using normal calculus operators like <code>+, -, *, /</code></p>
<p>e.g.:</p>
<pre><code class="lang-auto">f = theano.function([x], 2*x)
</code></pre>
<p>is equivuivalent to:</p>
<pre><code class="lang-auto"># suppose you have some tensor ``x``
# or create it
x = torch.zeros([100,100], device="cuda:0", dtype=torch.int)
def f(input):
    return input * 2
print(f(x))
</code></pre>
<p>Moreover, since pytorch is dynamic, sometimes users may require a just-in-time compilation utility to remove the tensor construction cost in python, you can do that by using <code>torch.jit</code></p>

username: ahmadreza9
cooked: <p>I know paradigm and your simple example which is explained in the first hyperlink of my comment. I want to create custom optimizer like this (<a href="https://discuss.pytorch.org/t/custom-optimizer-in-pytorch/22397" class="inline-onebox">Custom Optimizer in PyTorch</a>) so I should update weights like this:</p>
<pre><code class="lang-auto">weight_update = smth_with_good_dimensions
param.data.sub_(weight_update * learning_rate)
</code></pre>
<p>Now, how can we have a function in PyTorch like theano.function for my rms_prop update?</p>

username: Nikronic
cooked: <p>Hi,</p>
<p>As <a class="mention" href="/u/iffix">@iffiX</a> mentioned, we do not compile graph in PyTorch, it will be constructed during data flow in forward pass. So, for an optimizer, you just define a class that accepts parameters of a <code>model = nn.Module</code> and implements <code>step</code> function for it. Literally, based on your code, you need to only remove last line <code>self._train_fn = ...</code> as there is no compile stage in PyTorch.<br>
Here is the SGD implementation:<br>
<a href="https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD</a></p>
<p>Bests</p>

username: iffiX
cooked: <p>From the documentation of theano, I guess you are askinig for manual control over gradient updates?<br>
then:<br>
For simple modules, use <code>register_backward_hook</code><br>
Or control gradients on each of your input directly using: <code>register_hook</code></p>
<p>Optimizers are implemented in <code>torch.optim</code>, including RMSProp, Adam, SGD, etc.</p>

username: iffiX
cooked: <p>And indeed, if you want a completely new optimizer rather than control gradients by hand, you should inherit from <code>torch.optim.Optimizer</code> and implement the <code>step</code> method. (remember to wrap it with <span class="mention">@torch.no_grad</span>())</p>

username: ahmadreza9
cooked: <p>Thanks <a class="mention" href="/u/iffix">@iffiX</a>. This comment is the most relevant answer to my question. I will test  <code>register_backward_hook</code> function to manually control the gradient of my network(grads).</p>
<pre><code class="lang-auto">N = states.shape[0]

loss = T.log(prob_act[T.arange(N), actions]).dot(values) / N  # call it "loss"

grads = T.grad(loss, params)
</code></pre>