CNN predicts constant values for sparse amplitude regression — can't learn true pixel values

username: thunlok
cooked: <p>I’m training a small CNN (code below) to predict sparse amplitude maps from binary masks.</p>
<ul>
<li>
<p>Input: 60×60 image with exactly 15 pixels set to 1, rest are 0.</p>
</li>
<li>
<p>Target: Same size, 0 everywhere except those 15 pixels, which have values in the range 0.6–1.0.</p>
</li>
</ul>
<p>The CNN is trained on ~1800 images and tested on ~400. The goal is for it to predict the amplitude at the 15 known locations, given the mask as input.</p>
<p>Here’s an example output:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/3/a34222a0d2908564d4addbd58d278604e98dbfaa.png" data-download-href="https://discuss.pytorch.org/uploads/default/a34222a0d2908564d4addbd58d278604e98dbfaa" title="Screenshot from 2025-05-27 21-58-58"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/3/a34222a0d2908564d4addbd58d278604e98dbfaa_2_690x356.png" alt="Screenshot from 2025-05-27 21-58-58" data-base62-sha1="nifwlWD01nSVuzmpIridP8XrZrk" width="690" height="356" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/3/a34222a0d2908564d4addbd58d278604e98dbfaa_2_690x356.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/3/a34222a0d2908564d4addbd58d278604e98dbfaa_2_1035x534.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/3/a34222a0d2908564d4addbd58d278604e98dbfaa.png 2x" data-dominant-color="511460"><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 from 2025-05-27 21-58-58</span><span class="informations">1175×607 8.09 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>
<p>And the detailed predicted vs. target values:</p>
<pre><code>Index (row, col) |  Predicted |     Target
</code></pre>
<hr>
<pre><code>        (40, 72) |     0.9177 |     0.9143
        (40, 90) |     0.9177 |     1.0000
        (43, 52) |     0.9177 |     0.8967
        (50, 32) |     0.9177 |     0.9205
        (51, 70) |     0.9177 |     0.9601
        (53, 45) |     0.9177 |     0.9379
        (56, 88) |     0.9177 |     0.8906
        (61, 63) |     0.9177 |     0.9280
        (62, 50) |     0.9177 |     0.9154
        (65, 29) |     0.9177 |     0.9014
        (65, 91) |     0.9177 |     0.8941
        (68, 76) |     0.9177 |     0.9043
        (76, 80) |     0.9177 |     0.9206
        (80, 31) |     0.9177 |     0.8872
        (80, 61) |     0.9177 |     0.9019
</code></pre>
<p>As you can see, the network collapses to a constant output, despite the targets being quite different.<br>
I have been able to play around with the CNN and get values that are not all the same:</p>
<h2><a name="p-471102-index-row-col-predicted-target-1" class="anchor" href="#p-471102-index-row-col-predicted-target-1"></a>Index (row, col) |  Predicted |     Target</h2>
<pre><code>        (40, 72) |     0.9559 |     0.9143
        (40, 90) |     0.9563 |     1.0000
        (43, 52) |     0.9476 |     0.8967
        (50, 32) |     0.9515 |     0.9205
        (51, 70) |     0.9512 |     0.9601
        (53, 45) |     0.9573 |     0.9379
        (56, 88) |     0.9514 |     0.8906
        (61, 63) |     0.9604 |     0.9280
        (62, 50) |     0.9519 |     0.9154
        (65, 29) |     0.9607 |     0.9014
        (65, 91) |     0.9558 |     0.8941
        (68, 76) |     0.9560 |     0.9043
        (76, 80) |     0.9555 |     0.9206
        (80, 31) |     0.9620 |     0.8872
        (80, 61) |     0.9563 |     0.9019
</code></pre>
<p>I’ve tried many things:</p>
<ol>
<li>Scale the amplitudes to be from -5 to 5, -3 to 3, and -1 to 1 (linear and nonlinear behavior for them) then unscale when in the test() function</li>
<li>Different optimizers Adam and AdamW</li>
<li>Used different criteria: SmoothL1Loss() and MSELoss()</li>
<li>A large for loop over epoch and lr</li>
<li>Instead of doing a MSE for all pixels together, I instead did them individually</li>
</ol>
<p>What’s interesting is that I trained the same architecture for phase prediction, where values range from -π to π, and it learns beautifully:</p>
<pre><code>Index (row, col) |  Predicted |     Target
</code></pre>
<hr>
<pre><code>        (40, 72) |    -0.1235 |    -0.1235
        (40, 90) |     0.5146 |     0.5203
        (43, 52) |    -1.0479 |    -1.0490
        (50, 32) |    -0.3166 |    -0.3165
        (51, 70) |    -1.5540 |    -1.5521
        (53, 45) |     0.5990 |     0.6034
        (56, 88) |    -0.4752 |    -0.4752
        (61, 63) |    -2.4576 |    -2.4600
        (62, 50) |     2.0495 |     2.0526
        (65, 29) |    -2.6678 |    -2.6681
        (65, 91) |    -1.9935 |    -1.9961
        (68, 76) |    -1.9096 |    -1.9142
        (76, 80) |    -1.7976 |    -1.8025
        (80, 31) |    -2.7799 |    -2.7795
        (80, 61) |     0.5338 |     0.5393
</code></pre>
<p>So this proves that the CNN can learn, I just can’t figure out how it can work with amplitudes. The only difference is, that the input phase values are the same values as the loss function. Here is what I mean:</p>
<p>When being trained (let’s just take 1 pixel value of -1.2 for the phase):</p>
<p>-1.2 → CNN → output gets compared to -1.2</p>
<p>Whereas the amplitude of 1 pixel is like this:</p>
<p>1.0 → CNN -&gt;output gets compared to true value such as 0.9143</p>
<p>So maybe the phase has an “easier” life, nonetheless I am struggling with the CNN for the amplitude and I would really appreciate some insight if anyone can help!</p>
<pre><code>import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset
from sklearn.model_selection import train_test_split

class AmplitudeDataset(Dataset):
    def __init__(self, input_dir_ampl, target_dir_ampl):
        self.input_dir_ampl = input_dir_ampl
        self.target_dir_ampl = target_dir_ampl
        self.fileids = sorted([f.split('_')[1] for f in os.listdir(input_dir_ampl)])

    def __len__(self):
        return len(self.fileids)

    def __getitem__(self, idx):
        input_ampl = np.load(os.path.join(self.input_dir_ampl, f'tweezer_{self.fileids[idx]}_ampl_input.npy')).astype(np.float32)
        target_ampl = np.load(os.path.join(self.target_dir_ampl, f'tweezer_{self.fileids[idx]}_wgs_ampl.npy')).astype(np.float32)

        return torch.tensor(input_ampl).unsqueeze(0), torch.tensor(target_ampl).unsqueeze(0)

class SingleOutputFCN(nn.Module):
    def __init__(self):
        super(SingleOutputFCN, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding='same'),
            nn.LeakyReLU(0.1),
            nn.Conv2d(16, 16, kernel_size=3, padding='same'),
            nn.LeakyReLU(0.1),
            nn.Conv2d(16, 1, kernel_size=3, padding='same'),
        )

    def forward(self, x):
        return self.net(x)

def train(model, dataloader, device, epochs=4, lr=1e-3):
    model.to(device)
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    for epoch in range(epochs):
        total_loss = 0
        for inputs, targets in dataloader:
            inputs = inputs.to(device)
            targets = targets.to(device)

            optimizer.zero_grad()
            pred = model(inputs)
            loss = criterion(pred, targets)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            pred_std = pred.std().item()

        print(f"[AMPL] Epoch {epoch+1}/{epochs} | Avg Loss: {total_loss / len(dataloader):.6f}")

def test(model, dataloader, device):
    model.eval()
    output_dir = './cnn_predictions/'
    os.makedirs(output_dir, exist_ok=True)

    with torch.no_grad():
        for i, (inputs, _) in enumerate(dataloader):
            inputs = inputs.to(device)
            outputs = model(inputs).cpu()
            input_np = inputs.cpu().numpy()
            output_np = outputs.cpu().numpy()


            for j in range(inputs.shape[0]):
                idx = i * dataloader.batch_size + j
                file_id = dataloader.dataset.dataset.fileids[dataloader.dataset.indices[idx]] if isinstance(dataloader.dataset, Subset) else dataloader.dataset.fileids[idx]

                np.save(os.path.join(output_dir, f'tweezer_{file_id}_predicted_ampl.npy'), output_np[j, 0])
                np.save(os.path.join(output_dir, f'tweezer_{file_id}_input_ampl.npy'), input_np[j, 0])

if __name__ == '__main__':
    input_dir_ampl = './WGS_CNN_Cropped/ampl_input'
    target_dir_ampl = './WGS_CNN_Cropped/wgs_amplitude'

    dataset = AmplitudeDataset(input_dir_ampl, target_dir_ampl)
    train_ids, test_ids = train_test_split(range(len(dataset)), test_size=0.2, random_state=42)
    train_dataset = Subset(dataset, train_ids)
    test_dataset = Subset(dataset, test_ids)

    train_loader = DataLoader(train_dataset, batch_size=15, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=15, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = SingleOutputFCN()
    print("\n=== Training Amplitude Model ===")
    train(model, train_loader, device)

    print("\n=== Testing Amplitude Model ===")
    test(model, test_loader, device)
</code></pre>

username: KFrank
cooked: <p>Hi Thunlok!</p>
<aside class="quote no-group" data-username="thunlok" data-post="1" data-topic="220387" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/c4cdca/48.png" class="avatar"> thunlok:</div>
<blockquote>
<p>I’m training a small CNN (code below) to predict sparse amplitude maps from binary masks.</p>
<ul>
<li>
<p>Input: 60×60 image with exactly 15 pixels set to 1, rest are 0.</p>
</li>
<li>
<p>Target: Same size, 0 everywhere except those 15 pixels, which have values in the range 0.6–1.0.</p>
</li>
</ul>
</blockquote>
</aside>
<p>As I understand it, a given input sample is a 60x60 image that is mostly zero pixels and<br>
15 pixels that are 1.  Furthermore, the location of the 1 pixels differs from one input sample<br>
to the next.  So conceptually the problem is to predict the amplitude (and phase) for the<br>
pixels that are 1 in the input mask based solely on the locations of those 1 pixels.</p>
<aside class="quote no-group">
<blockquote>
<p>What’s interesting is that I trained the same architecture for phase prediction, where values range from -π to π, and it learns beautifully:</p>
<pre><code>Index (row, col) |  Predicted |     Target
</code></pre>
<hr>
<pre><code>        (40, 72) |    -0.1235 |    -0.1235
        (40, 90) |     0.5146 |     0.5203
        (43, 52) |    -1.0479 |    -1.0490
        (50, 32) |    -0.3166 |    -0.3165
        (51, 70) |    -1.5540 |    -1.5521
        (53, 45) |     0.5990 |     0.6034
        (56, 88) |    -0.4752 |    -0.4752
        (61, 63) |    -2.4576 |    -2.4600
        (62, 50) |     2.0495 |     2.0526
        (65, 29) |    -2.6678 |    -2.6681
        (65, 91) |    -1.9935 |    -1.9961
        (68, 76) |    -1.9096 |    -1.9142
        (76, 80) |    -1.7976 |    -1.8025
        (80, 31) |    -2.7799 |    -2.7795
        (80, 61) |     0.5338 |     0.5393
</code></pre>
<p>…</p>
<pre><code>class SingleOutputFCN(nn.Module):
    def __init__(self):
        super(SingleOutputFCN, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding='same'),
            nn.LeakyReLU(0.1),
            nn.Conv2d(16, 16, kernel_size=3, padding='same'),
            nn.LeakyReLU(0.1),
            nn.Conv2d(16, 1, kernel_size=3, padding='same'),
        )
</code></pre>
</blockquote>
</aside>
<p>This is very confusing to me as I don’t see how your model could possibly predict the<br>
phases.</p>
<p>If I understand your code correctly, your model pumps the input image through three<br>
convolutions with kernel size 3.  After the first 3x3 convolution, a given pixel that has<br>
value 1 will only affect a 3x3 square of surrounding pixels.  After two such convolutions,<br>
the affected square is 5x5, and after the third, it is 7x7.</p>
<p>Saying the same thing in a slightly different way, the prediction for a given output pixel<br>
can only depend on input pixels that are within three steps of the output-pixel location<br>
in both the x and y directions.</p>
<p>For the specific example you posted, I see pairs of input pixels that are no closer to one<br>
another (in both the x and y directions) than 8.</p>
<p>That is to say, the predicted output at the location of a given input pixel depends only on a<br>
7x7 field of 0 pixels with a single 1 pixel in the middle.  (This is because no other 1 pixel in<br>
the input image is close enough to fall within that 7x7 field.)  No information about the pattern<br>
and relative locations of the 1 pixels in the input ever makes it to the prediction for a specific<br>
output pixel.</p>
<p>There’s something I must be missing here, because the above analysis tells me that the<br>
predictions for inputs that look like the sample you posted – 1 pixels for which the separation<br>
in either the x or y direction is at least 8 – should just be a translationally-invariant constant<br>
that is about equal to the average of the target values.  (For the amplitudes you posted, the<br>
constant prediction is <code>0.9177</code> and the average target is <code>0.9182</code>)</p>
<p>Can you clarify what is going on here?  Are your input images truly binary – all the pixels are<br>
either exactly 0 or 1 (with 15 1 pixels in an input image)?</p>
<p>Best.</p>
<p>K. Frank</p>

username: thunlok
cooked: <p>Hi Krank,</p>
<p>Firstly, many thanks for taking the time to reply. I really do appreciate it.</p>
<p>I think your interpretation is spot on given the code structure I originally shared.</p>
<p>Let me clarify the intent and what I’m actually trying to achieve:</p>
<p>Yes, you’re correct that each input is a 60×60 binary mask with exactly 15 pixels set to 1 (and the rest 0), and those 1-pixel positions vary per sample. The target is a corresponding 60×60 array where the same 15 pixel locations have values between 0.6 and 1.0, representing amplitudes, and all other pixels are 0. So, in principle, the model’s task is to predict which value to assign at each active pixel location, based solely on the global arrangement of those 15 pixels.</p>
<p>You’re also absolutely right that in my initial CNN architecture the receptive field was too small to capture anything beyond the immediate 7×7 neighborhood. So essentially, the model was only learning to recognize the presence of a 1 pixel and defaulting to predicting the mean amplitude for all of them which is why I saw the constant collapse around ~0.9177.</p>
<p>However, the goal isn’t to predict amplitude purely from the local pattern around a 1 pixel it’s to infer amplitude from the global configuration of all 15 active pixels. The assumption is that the amplitude at a given location depends on where the other 14 are placed, possibly due to interference or overlap effects (eg in a holography context).</p>
<p>Does that make more sense? I think my CNN is too crude to do what I want it to do…I think something like using U-Net could work but I don’t know enough about it to implement it at this point.</p>
<p>Did that help? If not, let me know and I can happily re-explain.</p>
<p>Sincerely,<br>
thunlok</p>

username: thunlok
cooked: <p>Also I want to mention, I think unet, generative adversarial network, or graph neural network would be applicable here…but I don’t know enough about them to implement them</p>

username: KFrank
cooked: <p>Hi Thunlok!</p>
<aside class="quote no-group" data-username="thunlok" data-post="3" data-topic="220387" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/c4cdca/48.png" class="avatar"> thunlok:</div>
<blockquote>
<p>the model’s task is to predict which value to assign at each active pixel location, based solely on the global arrangement of those 15 pixels.</p>
</blockquote>
</aside>
<p>Yes, this is what I understood the use case to be.</p>
<aside class="quote no-group">
<blockquote>
<p>the model was only learning to recognize the presence of a 1 pixel and defaulting to predicting the mean amplitude for all of them which is why I saw the constant collapse around ~0.9177.</p>
</blockquote>
</aside>
<p>I’m still baffled as to why you didn’t get the same result when training on phases.  I would have<br>
expected (as with the amplitudes) the phase prediction to be a translationally-invariant constant.</p>
<aside class="quote no-group">
<blockquote>
<p>The assumption is that the amplitude at a given location depends on where the other 14 are placed, possibly due to interference or overlap effects (eg in a holography context).</p>
<p>Does that make more sense?</p>
</blockquote>
</aside>
<p>Yes.</p>
<aside class="quote no-group">
<blockquote>
<p>I think something like using U-Net could work.</p>
</blockquote>
</aside>
<p>U-Net also occurred to me as a possible approach.  You would want the receptive field of your<br>
U-Net, of course, to be large enough to contain all of the relevant active pixels, perhaps simply<br>
by having the receptive field contain the entire 60x60 input image.</p>
<p>Another approach – which to my mind might be more promising – would be to use as input<br>
to your model the coordinates of the active pixels.  That is, a single input sample would be<br>
a tensor of shape <code>[15, 2]</code> – the x and y coordinates of the 15 active pixels.  The model could<br>
then be a series of fully-connected layers.</p>
<p>This scheme would seemingly be easier for the model – the model doesn’t have to learn how<br>
to locate the active pixels and doesn’t have to learn that the locations of the active pixels are<br>
the things that matter.</p>
<p>Regardless of what approach you take, I would recommend training on both the amplitudes<br>
and phases simultaneously – more shared information.  Internal to the model, the learned<br>
“features” that help predict the amplitudes are likely to be able to help predict the phases,<br>
as well.</p>
<p>So in the scheme where you input the coordinates of the 15 active pixels (with shape <code>[15, 2]</code>),<br>
the target and predicted output would be the amplitude and phase of each of the active pixels,<br>
again with shape <code>[15, 2]</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: thunlok
cooked: <p>Hi Kfrank,</p>
<p>Thanks for the help! Actually the phase part – the CNN is just learning the identity transformation as the input phase and the phase that are included in the loss function are the same. We can disregard this for now as this is actually not that useful and I would like to make an impactful and robust method for this.</p>
<p>To give you more insight, I would like to do something like this paper: <a href="https://arxiv.org/pdf/2401.06014" rel="noopener nofollow ugc">https://arxiv.org/pdf/2401.06014</a></p>
<p>They use a U-net, but I want to go above their paper and try to use graph or generative neural networks. Does that make sense? It would be for the same use case as in the paper.</p>
<p>Sincerely,<br>
Thunlok</p>

username: thunlok
cooked: <p>Also, I didn’t mention, but the CNN should learn to have as input an arbitrary amount of activated pixels. So, I am just training on a training set of 1800 only with 15 pixels just to show the simplest case, but my goal is to expand this once it works and to give a random number of activated pixels (up to probably 500) and random position of those pixels. And then the goal would be to say “Hey CNN, take this amplitude input (and maybe also phase) of these X pixels at these positions and give me the best estimate of what the true result is”. Does that make sense?</p>

username: KFrank
cooked: <p>Hi Thunlok!</p>
<aside class="quote no-group" data-username="thunlok" data-post="6" data-topic="220387" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/c4cdca/48.png" class="avatar"> thunlok:</div>
<blockquote>
<p>try to use graph or generative neural networks.</p>
</blockquote>
</aside>
<p>I hate to rule anything out, but I don’t see anything in (what I think is) your use case that<br>
would make a graph or generative model appropriate.</p>
<aside class="quote no-group" data-username="thunlok" data-post="7" data-topic="220387">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/c4cdca/48.png" class="avatar"> thunlok:</div>
<blockquote>
<p>“Hey CNN, take this amplitude input (and maybe also phase) of these X pixels at these positions and give me the best estimate of what the true result is”.</p>
</blockquote>
</aside>
<p>It would be helpful if you could give an overview of your actual real-world use case.</p>
<p>From the various things you’ve said, I’m now imagining that your input is the locations of<br>
the active pixels together with some sort of measure of the amplitude and phase and you<br>
want to predict some kind of cleaned-up, maybe noise-reduced amplitude and phase for<br>
those pixels.</p>
<p>It still seems plausible to me that inputting the x and y coordinates of the active pixels<br>
(rather than force the model to learn to locate them) could be a better way to go.</p>
<p>As I understand it, samples of your real data have varying numbers of active pixels.<br>
Even so, you could still specify some maximum number of active pixels, say 500, and<br>
pad your input with “dummy” pixels to get an input length of 500.  You might also input<br>
a flag field that indicates which pixels are active.</p>
<p>So your input might have shape <code>[500, 5]</code> with each pixel (active or dummy) having<br>
the values <code>x</code>, <code>y</code>, <code>amplitude</code>, <code>phase</code>, and <code>flag</code>.  (I would probably use <code>flag = 1</code><br>
for active pixels and <code>flag = 0</code> for dummy pixels.)</p>
<p>It is true that with such a scheme the network would need to learn to pick the active<br>
pixels out of the 500 total pixels, active and dummy, that were input.  But this doesn’t<br>
seem any harder than having the network learn to find the active pixels in a 60x60<br>
image that contains 3600 pixels.</p>
<p>One last thought:  If your problem is in some sense translationally invariant – that is, if<br>
the prediction for a set of three or four nearby active pixels in the upper left of the image<br>
would be the same <em>regardless of where other active pixels are in the image,</em> as it would<br>
for the same pattern of three or four pixels in the lower right – then a (fully) convolutional<br>
network would be a natural fit, as the convolutions have that translational invariance built<br>
in.  But if the prediction for active pixels in the upper right depends globally on the locations<br>
of all of the active pixels – not just their nearby neighbors – then a convolutional network<br>
wouldn’t necessary have any special benefit for your use case.</p>
<p>Best.</p>
<p>K. Frank</p>

username: thunlok
cooked: <p>Hi Krank,</p>
<p>Thank you for the reply. I think you missed my additional comment where I mention the paper that I am trying to improve on. I am doing the same thing as them in terms of use case, but I want to have a better solution which would most likely need a more advanced network. Let me know what you think of the paper!</p>

username: thunlok
cooked: <p>I made a new post that goes into better detail of the use case if you would like to take a gander!</p>