How to accomodate batches when computing,visualizing reconstruction?

weight_theta: <p>Hi,</p>
<p>I have been working on a reconstruction of a continious signal for a vision project. Reconstruction is perfect when I do not use any batch information with a singal of lengt 3000 and <code>shape [3000]</code> , but when I try to wrap the signal into a single batch with shape  <code>[1, 3000]</code> my reconstruction error increase and looks like as if values are scaled by a factor of 0.5. Obviously this should not occur. I have tried examining this in more detail for the past days, but was unable to resolve this issue and am turning to the community.</p>
<p>While I tried to investigate different modes of error like : 1) whether batching is accounted for when computing the reconsturction, 2) trying to look for the where the scaling could occur,  or 3)  whether we are indexing into the correct values, all lead to a dead end. As such please find my code, for reference below. Please let me know if you spot something that could cause or resolve my problem.</p>
<h1><a name="libaries-1" class="anchor" href="#libaries-1"></a>Libaries</h1>
<pre><code class="lang-auto">import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import math
from scipy import linalg as la
from scipy import signal
from scipy import special as ss
device = torch.device("cpu")
</code></pre>
<h1><a name="model-2" class="anchor" href="#model-2"></a>Model</h1>
<pre><code class="lang-auto">def transition(measure, N, **measure_args):
    # Laguerre (translated)
    if measure == 'lagt':
        b = measure_args.get('beta', 1.0)
        A = np.eye(N) / 2 - np.tril(np.ones((N, N)))
        B = b * np.ones((N, 1))
    # Legendre (translated)
    elif measure == 'legt':
        Q = np.arange(N, dtype=np.float64)
        R = (2*Q + 1) ** .5
        j, i = np.meshgrid(Q, Q)
        A = R[:, None] * np.where(i &lt; j, (-1.)**(i-j), 1) * R[None, :]
        B = R[:, None]
        A = -A
    # Legendre (scaled)
    elif measure == 'legs':
        q = np.arange(N, dtype=np.float64)
        col, row = np.meshgrid(q, q)
        r = 2 * q + 1
        M = -(np.where(row &gt;= col, r, 0) - np.diag(q))
        T = np.sqrt(np.diag(2 * q + 1))
        A = T @ M @ np.linalg.inv(T)
        B = np.diag(T)[:, None]
        B = B.copy() # Otherwise "UserWarning: given NumPY array is not writeable..." after torch.as_tensor(B)
    elif measure == 'fourier':
        freqs = np.arange(N//2)
        d = np.stack([np.zeros(N//2), freqs], axis=-1).reshape(-1)[1:]
        A = 2*np.pi*(-np.diag(d, 1) + np.diag(d, -1))
        B = np.zeros(N)
        B[0::2] = 2
        B[0] = 2**.5
        A = A - B[:, None] * B[None, :]
        # A = A - np.eye(N)
        B *= 2**.5
        B = B[:, None]

    return A, B

class HiPPOScale(nn.Module):
    """ Vanilla HiPPO-LegS model (scale invariant instead of time invariant) """
    def __init__(self, N, method='legs', max_length=1024, discretization='bilinear'):
        """
        max_length: maximum sequence length
        """
        super().__init__()
        self.N = N
        A, B = transition(method, N)
        B = B.squeeze(-1)
        A_stacked = np.empty((max_length, N, N), dtype=A.dtype)
        B_stacked = np.empty((max_length, N), dtype=B.dtype)
        for t in range(1, max_length + 1):
            At = A / t
            Bt = B / t
            if discretization == 'forward':
                A_stacked[t - 1] = np.eye(N) + At
                B_stacked[t - 1] = Bt
            elif discretization == 'backward':
                A_stacked[t - 1] = la.solve_triangular(np.eye(N) - At, np.eye(N), lower=True)
                B_stacked[t - 1] = la.solve_triangular(np.eye(N) - At, Bt, lower=True)
            elif discretization == 'bilinear':
                A_stacked[t - 1] = la.solve_triangular(np.eye(N) - At / 2, np.eye(N) + At / 2, lower=True)
                B_stacked[t - 1] = la.solve_triangular(np.eye(N) - At / 2, Bt, lower=True)
            else: # ZOH
                A_stacked[t - 1] = la.expm(A * (math.log(t + 1) - math.log(t)))
                B_stacked[t - 1] = la.solve_triangular(A, A_stacked[t - 1] @ B - B, lower=True)
        self.register_buffer('A_stacked', torch.Tensor(A_stacked)) # (max_length, N, N)
        self.register_buffer('B_stacked', torch.Tensor(B_stacked)) # (max_length, N)

        vals = np.linspace(0.0, 1.0, max_length)
        self.eval_matrix = torch.Tensor((B[:, None] * ss.eval_legendre(np.arange(N)[:, None], 2 * vals - 1)).T  )

    def forward(self, inputs, fast=False):
        """
        inputs : (length, ...)
        output : (length, ..., N) where N is the order of the HiPPO projection
        """

        L = inputs.shape[0]

        inputs = inputs.unsqueeze(-1)
        u = torch.transpose(inputs, 0, -2)
        u = u * self.B_stacked[:L]
        u = torch.transpose(u, 0, -2) # (length, ..., N)

        if fast:
            result = unroll.variable_unroll_matrix(self.A_stacked[:L], u)
            return result

        c = torch.zeros(u.shape[1:]).to(inputs)
        cs = []
        for t, f in enumerate(inputs):
            c = F.linear(c, self.A_stacked[t]) + self.B_stacked[t] * f
            cs.append(c)
        return torch.stack(cs, dim=0)

    def reconstruct(self, c):
        a = self.eval_matrix.to(c) @ c.unsqueeze(-1)
        return a
</code></pre>
<h1><a name="synthetic-data-generation-3" class="anchor" href="#synthetic-data-generation-3"></a>Synthetic data generation</h1>
<pre><code class="lang-auto">def whitesignal(period, dt, freq, rms=0.5, batch_shape=()):
    
    # Produces output signal of length period / dt, band-limited to frequency freq
    # Output shape (*batch_shape, period/dt)
     
    if freq is not None and freq &lt; 1. / period:
        raise ValueError(f"Make ``{freq=} &gt;= 1. / {period=}`` to produce a non-zero signal",)

    nyquist_cutoff = 0.5 / dt
    if freq &gt; nyquist_cutoff:
        raise ValueError(f"{freq} must not exceed the Nyquist frequency for the given dt ({nyquist_cutoff:0.3f})")

    n_coefficients = int(np.ceil(period / dt / 2.))
    shape = batch_shape + (n_coefficients + 1,)
    sigma = rms * np.sqrt(0.5)
    coefficients = 1j * np.random.normal(0., sigma, size=shape)
    coefficients[..., -1] = 0.
    coefficients += np.random.normal(0., sigma, size=shape)
    coefficients[..., 0] = 0.

    set_to_zero = np.fft.rfftfreq(2 * n_coefficients, d=dt) &gt; freq
    coefficients *= (1-set_to_zero)
    power_correction = np.sqrt(1. - np.sum(set_to_zero, dtype=float) / n_coefficients)
    if power_correction &gt; 0.: coefficients /= power_correction
    coefficients *= np.sqrt(2 * n_coefficients)
    signal = np.fft.irfft(coefficients, axis=-1)
    signal = signal - signal[..., :1]  # Start from 0
    return signal

</code></pre>
<h1><a name="reconstruction-4" class="anchor" href="#reconstruction-4"></a>Reconstruction</h1>
<pre><code class="lang-auto">def reconstruct(T, dt, N, freq, vals, u):
    u = torch.tensor(u, dtype=torch.float)
    u = u.to(device)

    # Linear Time Invariant (LTI) methods x' = Ax + Bu
    lti_methods = [
        'legt',
    ]

    # Original HiPPO-LegS, which uses time-varying SSM x' = 1/t [ Ax + Bu]
    # we call this "linear scale invariant"
    lsi_methods = ['legs']
    for method in lsi_methods:
        hippo = HiPPOScale(N=N, method=method, max_length=int(T/dt)).to(device)
        u_hippo = hippo.reconstruct(hippo(u))[-1].cpu()
        u_hippo_all = hippo.reconstruct(hippo(u)).cpu()
    
    return u, u_hippo, u_hippo_all, vals`
</code></pre>
<h1><a name="visualise-signal-with-no-batch-information-5" class="anchor" href="#visualise-signal-with-no-batch-information-5"></a>Visualise signal with no batch information</h1>
<pre><code class="lang-auto">u = whitesignal(3, 1e-3, 3.0, batch_shape=())
sig, recon, recon_all_t, values = reconstruct(T=3, dt=1e-3, N=64, freq=3.0, vals=vals, u=u)

plt.plot(values [-len(sig):], sig, label='u=target',linewidth=2, dashes=(5, 10))
plt.plot(values [-len(recon):], recon, label='recon', dashes=(5, 1))
plt.legend()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/b/8bae4039600855556f94f4d3f1a81224a2085a34.png" data-download-href="https://discuss.pytorch.org/uploads/default/8bae4039600855556f94f4d3f1a81224a2085a34" title="Screenshot_20230221_125350"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8bae4039600855556f94f4d3f1a81224a2085a34_2_671x500.png" alt="Screenshot_20230221_125350" data-base62-sha1="jVFHlC6E7ceWsDhTcEGMTCwF2mw" width="671" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/b/8bae4039600855556f94f4d3f1a81224a2085a34_2_671x500.png, https://discuss.pytorch.org/uploads/default/original/3X/8/b/8bae4039600855556f94f4d3f1a81224a2085a34.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/b/8bae4039600855556f94f4d3f1a81224a2085a34.png 2x" data-dominant-color="F6F6F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot_20230221_125350</span><span class="informations">980×730 126 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

weight_theta: <h1><a name="visualise-signal-with-1-batch-dimension-1" class="anchor" href="#visualise-signal-with-1-batch-dimension-1"></a>Visualise signal with 1 batch dimension</h1>
<pre><code class="lang-auto">plt.plot(values[-len(sig[0]):], sig[0], label='target',linewidth=3,)
plt.plot(values[-len(recon[:,-1,0]):], recon[:,-1,0], linestyle='--', label='recon', dashes=(5, 1))
plt.legend()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/9/d9d300e475bcaaa4e36233df5edd1ed4baae037d.png" data-download-href="https://discuss.pytorch.org/uploads/default/d9d300e475bcaaa4e36233df5edd1ed4baae037d" title="Screenshot_20230221_010432"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d9d300e475bcaaa4e36233df5edd1ed4baae037d_2_678x500.png" alt="Screenshot_20230221_010432" data-base62-sha1="v4XGpbvRlm9LSfelQYF2uG4im4Z" width="678" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d9d300e475bcaaa4e36233df5edd1ed4baae037d_2_678x500.png, https://discuss.pytorch.org/uploads/default/original/3X/d/9/d9d300e475bcaaa4e36233df5edd1ed4baae037d.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/9/d9d300e475bcaaa4e36233df5edd1ed4baae037d.png 2x" data-dominant-color="F7F7F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot_20230221_010432</span><span class="informations">975×719 120 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Based on your description I would guess your code might be accidentally broadcast some tensors resulting in a wrong computation. Could you check the shapes of the intermediate tensors and see if that’s the case?</p>

weight_theta: <p>Thank you for you answer ptrblck. I have looked into shapes of the intermediate variables and added them for reference below. What would I be interested in ? I thought I have is that the coefficent matrix may not be correctly initalised, shouldnt it be <code>[batchdim, N] [1, 64]</code> instead of <code>[3000, 64]</code> ? This then propogates through to when <code>Coefs c</code> are computed and similarly should be <code>[1, 64]</code> instead of<code> [3000, 64]</code>. But I am not entirely sure as the other matrices like <code>Coefs stacked shape</code> looks correct,<code> Eval mat</code> also looks correct and the recon shape as well. What do you think Patrick ?</p>
<h1>
<a name="shapes-intermediate-variables-no-batch-information-1" class="anchor" href="#shapes-intermediate-variables-no-batch-information-1"></a>Shapes intermediate variables no batch information</h1>
<pre><code class="lang-auto">-------------Init------------ torch.Size([64])
Coef c shape: torch.Size([64])
-------------Init------------ torch.Size([64])
Coef  c shape: torch.Size([64])
----------------------------------------
----------Reconstruction part:----------
Eval mat shape: torch.Size([3000, 64])
Coefs stacked shape: torch.Size([3000, 64])
Recon shape: torch.Size([3000, 3000, 1])
----------------------------------------
-------------Init------------ torch.Size([64])
Coef c shape: torch.Size([64])
-------------Init------------ torch.Size([64])
Coef  c shape: torch.Size([64])
----------------------------------------
----------Reconstruction part:----------
Eval mat shape: torch.Size([3000, 64])
Coefs stacked shape: torch.Size([3000, 64])
Recon shape: torch.Size([3000, 3000, 1])
----------------------------------------
</code></pre>
<h1>
<a name="shapes-intermediate-variables-1-batch-dimension-2" class="anchor" href="#shapes-intermediate-variables-1-batch-dimension-2"></a>Shapes intermediate variables 1 batch dimension</h1>
<pre><code class="lang-auto">-------------Init------------ torch.Size([3000, 64])
Coef c shape: torch.Size([3000, 64])
-------------Init------------ torch.Size([3000, 64])
Coef  c shape: torch.Size([3000, 64])
----------------------------------------
----------Reconstruction part:----------
Eval mat shape: torch.Size([3000, 64])
Coefs stacked shape: torch.Size([1, 3000, 64])
Recon shape: torch.Size([1, 3000, 3000, 1])
----------------------------------------
-------------Init------------ torch.Size([3000, 64])
Coef c shape: torch.Size([3000, 64])
-------------Init------------ torch.Size([3000, 64])
Coef  c shape: torch.Size([3000, 64])
----------------------------------------
----------Reconstruction part:----------
Eval mat shape: torch.Size([3000, 64])
Coefs stacked shape: torch.Size([1, 3000, 64])
Recon shape: torch.Size([1, 3000, 3000, 1])
----------------------------------------
</code></pre>

ptrblck: <p>Thanks for the update!<br>
Some shapes look indeed incorrect and you might need to check these, but let me give you a quick overview first.<br>
In PyTorch most layers expect an input tensor in the shape <code>[batch_size, *]</code>, where <code>*</code> denotes additional dimensions. There are a few exceptions, such as RNNs, which expect the batch dimension in <code>dim1</code> by default (you can change it via <code>bathc_first=True</code>), and the latest PyTorch release also accepts inputs without the batch dimension (it will <code>unsqueeze</code> the missing batch dimension for you), but let’s stick to the common approach.<br>
All intermediate acitvations will thus also have a shape of <code>[batch_size, *]</code>, but note that the trainable parameters and buffers do not depend on the batch size and thus do not have a batch dimension.<br>
Take a look at this simple example:</p>
<pre><code class="lang-python">lin = nn.Linear(8, 16)
print(lin.weight.shape)
# torch.Size([16, 8])
print(lin.bias.shape)
# torch.Size([16])

x = torch.randn(1, 8)
out = lin(x)
print(out.shape)
# torch.Size([1, 16])

x = torch.randn(20, 8)
out = lin(x)
print(out.shape)
# torch.Size([20, 16])
</code></pre>
<p>You can see that the <code>weight</code> and <code>bias</code> parameters do not care about the actual batch size and work with different inputs.</p>
<p>In your posted output a few tensors seem to change their shape if a batch size of <code>1</code> was used for the input:</p>
<pre><code class="lang-python"># from
Coef c shape: torch.Size([64])
# to
Coef c shape: torch.Size([3000, 64])
</code></pre>
<p>which seems unexpected.<br>
I don’t know which exact parameter this is, but even if this would be a forward activation I would expect to see a new dimension with a size of <code>1</code> (<code>batch_size</code>) instead of <code>3000</code>.<br>
Could you check which parameter this is and how its shape depends on the batch size?</p>