Computationally similar functions result in different behavior. Am I missing someting wrt to autograd

username: Paarulakan
cooked: <p>I have to refactor a old physics informed neural network code for further experiments. The old code is in notebook form and has lot of explicit computational expressions. As you can see in the new implementation, it is much more flexible which is necessary for our work right now. The following are two snippets of a function that I found out to be the culprit. I have replaced all the old code with new code incrementally and test out by training the model and comparing metrics and visualizing the functions modeled by the PINN.  All of my new code changes have no unintended effect and reproduce the results of old notebook with error differences in the order of 1e-5.</p>
<p>Only when I use the new net_f function as defined in the training the training gets slow by a large measure. For instance with old notebook’s net_f by 1900 iterations, the error reached 1e-03 range where as using new net_f definition below the error is still in 1e-1 range.</p>
<p>Both implementation use the LBFGS optimizer with following config</p>
<pre data-code-wrap="python"><code class="lang-python">    # optimizers: using the same settings
          self.optimizer = torch.optim.LBFGS(
              self.dnn.parameters(),
              lr=1.0,
              max_iter=50000,
              max_eval=50000,
              history_size=50,
              tolerance_grad=1e-5,
              tolerance_change=1.0 * np.finfo(float).eps,
              line_search_fn="strong_wolfe"       # can be "strong_wolfe"
          )

</code></pre>
<p><strong>Note:</strong> The names is just string <code>uvpsab</code></p>
<pre data-code-wrap="python"><code class="lang-python">   def cleave_and_name(self, input_tensor):
       return OrderedDict([(name, input_tensor[:, i:j])
                           for name, i, j in zip(self.names,
                                                 range(NCOMP),
                                                 range(1, NCOMP+1))])
   def name_the_tensors(self, tensors):
       return OrderedDict([(name, i) for name, i in zip(self.names, tensors)])

   def net_f(self, x, t):
       """ The pytorch autograd version of calculating residual """
       y = self.net_u(x, t)

       F = self.cleave_and_name(y).values()

       def compute_grad(diff, wrt):
           return torch.autograd.grad(diff, wrt,
                                      grad_outputs=torch.ones_like(diff),
                                      retain_graph=True,
                                      create_graph=True)[0]

       Ft = [compute_grad(i, t) for i in F]
       Fx = [compute_grad(i, x) for i in F]
       Fxx = [compute_grad(i, x) for i in Fx]

       def compute_output(f, f_t, f_xx, c1, c2, c1_sign):
           return f_t + (c1_sign * c1 * f_xx) + ((-1 * c1_sign) * c2 * f)

       N_F = len(F)
       SQ = sum([ i**2 for i in F ])
       C1 = [0.5] * N_F
       C2 = [SQ] * N_F
       C1_SIGN = [pow(-1, i) for i in range(1, N_F+1)]

       F_output = [compute_output(*i)
                   for i in zip(F, Ft, Fxx, C1, C2, C1_SIGN)]

       return self.name_the_tensors(F_output)
</code></pre>
<p>The old implementation is as follows.</p>
<pre data-code-wrap="python"><code class="lang-python">    def net_f(self, x, t):
        """ The pytorch autograd version of calculating residual """
        y = self.net_u(x, t)

        u = y[:, 0:1]
        v = y[:, 1:2]
        p = y[:, 2:3]
        s = y[:, 3:4]
        a = y[:, 4:5]
        b = y[:, 5:6]

        u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        v_t = torch.autograd.grad(v, t, grad_outputs=torch.ones_like(v), retain_graph=True, create_graph=True)[0]

        p_t = torch.autograd.grad(p, t, grad_outputs=torch.ones_like(p), retain_graph=True, create_graph=True)[0]
        s_t = torch.autograd.grad(s, t, grad_outputs=torch.ones_like(s), retain_graph=True, create_graph=True)[0]

        a_t = torch.autograd.grad(a, t, grad_outputs=torch.ones_like(a), retain_graph=True, create_graph=True)[0]
        b_t = torch.autograd.grad(b, t, grad_outputs=torch.ones_like(b), retain_graph=True, create_graph=True)[0]

        u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), retain_graph=True, create_graph=True)[0]
        v_x = torch.autograd.grad(v, x, grad_outputs=torch.ones_like(v), retain_graph=True, create_graph=True)[0]

        p_x = torch.autograd.grad(p, x, grad_outputs=torch.ones_like(p), retain_graph=True, create_graph=True)[0]
        s_x = torch.autograd.grad(s, x, grad_outputs=torch.ones_like(s), retain_graph=True, create_graph=True)[0]

        a_x = torch.autograd.grad(a, x, grad_outputs=torch.ones_like(a), retain_graph=True, create_graph=True)[0]
        b_x = torch.autograd.grad(b, x, grad_outputs=torch.ones_like(b), retain_graph=True, create_graph=True)[0]


        u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), retain_graph=True, create_graph=True)[0]
        v_xx = torch.autograd.grad(v_x, x, grad_outputs=torch.ones_like(v_x), retain_graph=True, create_graph=True)[0]

        p_xx = torch.autograd.grad(p_x, x, grad_outputs=torch.ones_like(p_x), retain_graph=True, create_graph=True)[0]
        s_xx = torch.autograd.grad(s_x, x, grad_outputs=torch.ones_like(s_x), retain_graph=True, create_graph=True)[0]

        a_xx = torch.autograd.grad(a_x, x, grad_outputs=torch.ones_like(a_x), retain_graph=True, create_graph=True)[0]
        b_xx = torch.autograd.grad(b_x, x, grad_outputs=torch.ones_like(b_x), retain_graph=True, create_graph=True)[0]


        SQ = u**2 + v**2 + p**2 + s**2 + a**2 + b**2

        f_u = v_t - 0.5*u_xx + SQ*u

        f_v = u_t + 0.5*v_xx - SQ*v

        f_p = s_t - 0.5*p_xx + SQ*p

        f_s = p_t + 0.5*s_xx - SQ*s

        f_a = b_t - 0.5*a_xx + SQ*a

        f_b = a_t + 0.5*b_xx - SQ*b


        return f_u, f_v, f_p, f_s, f_a, f_b

</code></pre>

username: KFrank
cooked: <p>Hi Paarulakan!</p>
<p>It’s pretty hard to understand what you’re doing here.</p>
<aside class="quote no-group" data-username="Paarulakan" data-post="1" data-topic="208950" data-full="true">
<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/paarulakan/48/13733_2.png" class="avatar"> Paarulakan:</div>
<blockquote>
<p>All of my new code changes have no unintended effect</p>
</blockquote>
</aside>
<p>Do all of your “new code changes” include your new “net_f function?”</p>
<aside class="quote no-group">
<blockquote>
<p>reproduce the results of old notebook with error differences in the order of 1e-5.</p>
</blockquote>
</aside>
<p>What is it that doesn’t change by more than 1e-5?  One forward pass?<br>
One backward pass?  The result of an optimization step?</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">              self.dnn.parameters(),
</code></pre>
</blockquote>
</aside>
<p>What is <code>self.dnn</code>?</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">        y = self.net_u(x, t)
</code></pre>
</blockquote>
</aside>
<p>Is <code>self.dnn</code> the same as <code>self.net_u</code>?</p>
<p><code>LBFGS.step()</code> takes a closure that, among other things, returns<br>
a scalar loss value.  You haven’t showed us your closure nor your<br>
call to <code>LBFGS.step()</code>, so we have no idea where things might be<br>
going wrong.</p>
<p>The differences of 1e-5 you mention would seem to be consistent<br>
with single-precision round-off error.  But, in any event, you should<br>
try repeating your computation in double precision to make sure<br>
that your new vs. old code isn’t introducing some error not due to<br>
round-off error.</p>
<p>Presumably you compute a scalar loss value somewhere.  Try<br>
performing a single forward pass with both your old and new <code>net_f</code>.<br>
Do the two computations of your loss value agree up to round-off<br>
error?</p>
<p>It’s not at all clear which parameters you are optimizing and with<br>
respect to which parameters you would be computing gradients.<br>
But, whatever they are, try performing a single backward pass<br>
with both your old and new <code>net_f</code>.  Do the resulting gradients<br>
agree up to round-off error?</p>
<p>Then try performing a single optimization step, and so on.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: Paarulakan
cooked: <p>First of all thank you so much for going to the trouble of reading this and responding kindly.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="208950">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Do all of your “new code changes” include your new “net_f function?”</p>
</blockquote>
</aside>
<p>Yes. So this is how it went. There is original notebook that produces a set of plot. I refactored the notebook to be more generic/easily adaptable which includes everything from the original notebook. But when I ran my new code the result were not replicated. The training was orders of magnitude slow i.e the loss trend (decrease per iteration) did not follow the original notebook’s loss trend. So replaced code blocks in my new code with code blocks from original notebook and tested for loss trend. everything in my new code works except net_f. All of my new implementations of the original functions act computationally same except for net_f.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="208950">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>What is it that doesn’t change by more than 1e-5? One forward pass?<br>
One backward pass? The result of an optimization step?</p>
</blockquote>
</aside>
<p>it is the difference in error value of new code(with old impl of net_f) vs old code after training occurred for 10k iterations. So net_f is the culprit. I have provided both the implementations of net_f in the OP(top one new, bottom one old).</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="208950">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Is <code>self.dnn</code> the same as <code>self.net_u</code>?</p>
</blockquote>
</aside>
<p>Pretty much. self.net_u() is also a function which just concats [x,t] and runs self.dnn on that.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="208950">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p><code>LBFGS.step()</code> takes a closure that, among other things, returns<br>
a scalar loss value. You haven’t showed us your closure nor your<br>
call to <code>LBFGS.step()</code>, so we have no idea where things might be<br>
going wrong.</p>
<p>The differences of 1e-5 you mention would seem to be consistent<br>
with single-precision round-off error. But, in any event, you should<br>
try repeating your computation in double precision to make sure<br>
that your new vs. old code isn’t introducing some error not due to<br>
round-off error.</p>
<p>Presumably you compute a scalar loss value somewhere. Try<br>
performing a single forward pass with both your old and new <code>net_f</code>.<br>
Do the two computations of your loss value agree up to round-off<br>
error?</p>
</blockquote>
</aside>
<p>Yes, there is function named loss_func() which is the closure supplied to optimizer.step(). loss_func runs both net_u and net_f and performs gradient descent using mean squared error. the new impl of loss_func is shown below. This works. I confirmed it by running for 10k iterations and results are replicated. If you want old impl of loss_func I can share the code, but it is too lengthy and basically inlined variant of the new loss_func similar to old net_f impl.</p>
<pre data-code-wrap="python"><code class="lang-python">def loss_func(self):
        self.iter += 1
        self.optimizer.zero_grad()

        pred_icbc = self.net_u(self.x_u, self.t_u)
        pred_f = self.net_f(self.x_f, self.t_f)

        pred_icbc = self.cleave_and_name(pred_icbc)
        loss_icbc = OrderedDict([
            (name, torch.mean((self.icbc[name] - pred_icbc[name])**2)) for name in self.names
        ])

        loss_f = OrderedDict([
            (name, torch.mean(pred_f[name]**2)) for name in self.names
        ])

        loss = sum(loss_icbc.values()) + sum(loss_f.values())
        loss.backward()

        if self.iter % self.checkpoint == 0:
            elapsed = time.time() - self.training_start_time
            print(f'iteration: {self.iter} / elapsed: {elapsed} / loss: {loss.item():3e}')
            print('  icbc loss:' + '/'.join([f'{n}: {l.item():3e}' for n, l in loss_icbc.items()]))
            print('  f loss:' + '/'.join([f'{n}: {l.item():3e}' for n, l in loss_f.items()]))

            torch.save({
                'iter': self.iter,
                'model_state_dict': self.dnn.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
                'loss': loss,
            },'model.pt')

        return loss
</code></pre>
<p>The thing that puzzles me is that the both net_f functions are semantically same and should be computationally same too which they are evidently not. I thought my lack of understanding autograd might be a problem.</p>

username: KFrank
cooked: <p>Hi Paarulakan!</p>
<aside class="quote no-group" data-username="Paarulakan" data-post="3" data-topic="208950" data-full="true">
<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/paarulakan/48/13733_2.png" class="avatar"> Paarulakan:</div>
<blockquote>
<p>it is the difference in error value of new code(with old impl of net_f) vs old code after training occurred for 10k iterations.</p>
</blockquote>
</aside>
<p>In general, trying to debug by observing some numerical discrepancy<br>
after 10k iterations is a bad approach.  All of those iterations permit<br>
round-off error to accumulate, mix things together, and generally<br>
confuse the issue.</p>
<aside class="quote no-group">
<blockquote>
<p>there is function named loss_func() which is the closure supplied to optimizer.step().<br>
…</p>
<pre data-code-wrap="python"><code class="lang-python">def loss_func(self):
        ....
        loss = sum(loss_icbc.values()) + sum(loss_f.values())
        loss.backward()
        ....
        return loss
</code></pre>
</blockquote>
</aside>
<p>This makes sense.  <code>loss_func()</code> (the <code>LBFGS</code> closure) performs<br>
one forward pass followed by a backward pass and returns the<br>
loss value.</p>
<p>Start with this (leaving <code>LBFGS</code> out of it).  First call <code>loss_func()</code><br>
once with your old and new <code>net_f()</code> functions and compare both<br>
the loss value and the computed gradients (that I assume get<br>
populated into the <code>.grad</code> properties of <code>self.dnn</code>’s <code>Parameter</code>s).</p>
<p>Do these values agree up to some reasonable round-off error (after<br>
just a single forward / backward pass, i.e, after just a single call to<br>
<code>loss_func()</code>)?  Do this in both single and double precision just to<br>
make sure that you’re not fooling yourself.</p>
<p>If they do agree, perform a single call to <code>LBFGS.step (loss_func)</code><br>
and compare the updated <code>Parameter</code> values (presumably those of<br>
<code>self.dnn</code>).  Again, try doing this in both single and double precision<br>
to make sure that a “real” discrepancy isn’t hiding itself behind what<br>
might look like round-off error.</p>
<p>The basic idea is to perform divide-and-conquer debugging where<br>
you try to locate where the error first occurs (or at least where the<br>
discrepancy can first be detected).  Also, don’t focus on “derived”<br>
quantities like how your training is converging, but instead on more<br>
basic quantities like the loss value, the gradients, or the updated<br>
parameter values.</p>
<p>If you don’t see a discrepancy after a single <code>LBFGS.step()</code>, try<br>
running (smallish) numbers of iterations until you do.</p>
<p>Conversely, if you do see a discrepancy (but can’t track down its<br>
cause), try running plain-vanilla <code>SGD</code> as your optimizer where you<br>
can run a single step as a time (<code>LBFGS</code> takes multiple steps under<br>
the hood), and then increase the number of <code>SGD</code> iterations until<br>
you can detect a discrepancy.</p>
<p>Assuming that you’re right about <code>net_f()</code> being the culprit, after<br>
you’ve located the first detectable discrepancy, you can now print<br>
out intermediate values in both the old and new <code>net_f()</code>s to nail<br>
down which specific operation differs between the two versions.</p>
<p>Good luck!</p>
<p>K. Frank</p>