CNN and noise filtering

DiWarp: <p>Hello guys, hope you are all alright.</p>
<p>I’m facing a problem here. I’m trying to implement a ResNet18 to be able to filter noise. For that, I’m using the MNIST dataset and injecting noise on it. The code runs fine, no error message is shown, but here comes the problem:</p>
<p>the code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torchvision import models
import matplotlib.pyplot as plt
from PIL import Image

# class NoisyMNIST
class NoisyMNIST(MNIST):
    def __init__(self, root, train=True, transform=None, noisy_transform=None, download=False):
        super().__init__(root, train=train, transform=transform, download=download)
        self.noisy_transform = noisy_transform

    def __getitem__(self, index):
        img = self.data[index].view(28, 28)  # Remove the normalization and conversion to tensor

        if self.transform is not None:
            clean_img = self.transform(img)

        if self.noisy_transform is not None:
            noisy_img = self.noisy_transform(img)

        return noisy_img, clean_img


#pre-trained resnet
resnet = models.resnet18(pretrained=True)


resnet = nn.Sequential(*list(resnet.children())[:-1])


resnet.add_module('upsample', nn.Upsample(size=(28, 28), mode='bilinear', align_corners=False))
resnet.add_module('leakyrelu', nn.LeakyReLU())
resnet.add_module('dropout', nn.Dropout(p=0.5))
resnet.add_module('last_conv', nn.Conv2d(512, 3, kernel_size=1))
torch.nn.init.kaiming_normal_(resnet.last_conv.weight)

# loss and optimizer
criterion = nn.L1Loss()
optimizer = torch.optim.Adam(resnet.parameters(), lr=0.001)

# Load dataset and add gaussian noise to the images
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.float()),  # Convert to FloatTensor
    transforms.Normalize((0.5,), (0.5,)),
    transforms.Lambda(lambda x: x.repeat(3, 1, 1))  # repeating channel
])

noisy_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.float()),  # Convert to FloatTensor
    transforms.Normalize((0.5,), (0.5,)),
    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),  # repeating channel
    transforms.Lambda(lambda x: x + torch.randn_like(x))  # adding noise
])


train_set = NoisyMNIST('./data', train=True, transform=transform, noisy_transform=noisy_transform, download=True)
test_set = NoisyMNIST('./data', train=False, transform=transform, noisy_transform=noisy_transform, download=True)

# Split MNIST dataset
# train_set, valid_set = torch.utils.data.random_split(train_set, [50000, 10000])

# Reducing it to 10%
num_train = int(len(train_set) * 0.1)
num_valid = int(len(test_set) * 0.05)  # Usamos metade para validação e metade para teste
num_test = num_valid

train_set, _ = random_split(train_set, [num_train, len(train_set) - num_train])
valid_set, _ = random_split(test_set, [num_valid, len(test_set) - num_valid])
test_set, _ = random_split(test_set, [num_test, len(test_set) - num_test])

train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=64, shuffle=False)
test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

# To train the model
num_epochs = 10
train_losses = []
valid_losses = []

for i, data in enumerate(train_loader, 0):
    inputs, targets = data
    inputs = inputs.float()
    targets = targets.float()

    print(f"inputs shape: {inputs.shape}")
    print(f"targets shape: {targets.shape}")

for epoch in range(num_epochs):
    print(f"Training epoch {epoch+1}/{num_epochs}...")
    running_loss = 0.0

    # Training
    for i, data in enumerate(train_loader, 0):
        inputs, targets = data
        inputs = inputs.float()
        targets = targets.float()

        optimizer.zero_grad()
        outputs = resnet(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    # Validating
    with torch.no_grad():
        for i, data in enumerate(valid_loader, 0):
            inputs, targets = data
            inputs = inputs.float()
            targets = targets.float()

            outputs = resnet(inputs)
            loss = criterion(outputs, targets)

            valid_losses.append(loss.item())

    train_losses.append(running_loss/len(train_loader))

print("Training completed!")

# Evaluate model
print("Evaluating model...")
test_losses = []
with torch.no_grad():
    for i, data in enumerate(test_loader, 0):
        inputs, targets = data
        inputs = inputs.float()
        targets = targets.float()

        outputs = resnet(inputs)
        loss = criterion(outputs, targets)

        test_losses.append(loss.item())

# Plot loss 
plt.plot(train_losses, label='Treino')
plt.plot(valid_losses, label='Validação')
plt.legend()
plt.title('Perda')
plt.show()

# Save the trained model
#torch.save(resnet.state_dict(), './resnet_noise_removal.pth')

# Show the input images and the denoised images
plt.figure(figsize=(10, 5))
for i in range(5):
    plt.subplot(2, 5, i+1)
    plt.imshow(inputs[i].permute(1, 2, 0).numpy(), cmap='gray')
    plt.title('Input')
    plt.axis('off')

    plt.subplot(2, 5, i+6)
    plt.imshow(outputs[i].detach().permute(1, 2, 0).numpy(), cmap='gray')
    plt.title('Denoised')
    plt.axis('off')
plt.show()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/e/2e3022db200d9f4356c15ec0ec7458c5ad5546c6.png" alt="progress" data-base62-sha1="6AB5SIqnIxhUzPcL0AXZ53Qmo1U" width="278" height="218"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/4/f41a5770830b89b5323e053b431d5fe2ec1ff9bf.png" data-download-href="https://discuss.pytorch.org/uploads/default/f41a5770830b89b5323e053b431d5fe2ec1ff9bf" title="download"><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/4/f41a5770830b89b5323e053b431d5fe2ec1ff9bf.png" alt="download" data-base62-sha1="yPqWeKUDn0nIeR3ZhnusBIhYgPR" width="690" height="334" data-dominant-color="727373"><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">download</span><span class="informations">795×385 9.77 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>It’s not returning the image without noise, it’s returning all black.</p>
<p>What am I doing wrong? What can I do to fix this?</p>
<p>Thanks!</p>

ptrblck: <p>Try to overfit a tiny subset of your data (e.g. just 10 samples) by playing around with the hyperparameters of your training. Once your model is able to overfit the data you could try to scale up the use case again.</p>

dhruvbird: <p>Here are some things I would try (can’t guarantee that any of them will work though).</p>
<ol>
<li>Drop the normalization of (0.5,), (0.5,) since that’s not how the data was normalized when the pre-trained model was trained. Just try w/o any normalization for now. Additionally, the model was trained on colour images so the input distribution is probably completely different.</li>
<li>Don’t use a pre-trained model. For this specific use-case of MNIST images, you probably don’t need a pre-trained model - and it’s probably hurting you.</li>
<li>Check the range of output values. Ideally, you want them to be in the range [0.0 … 1.0]. If they aren’t (which is likely) then you can force it by applying a <code>sigmoid()</code> on the output tensor.</li>
</ol>
<p>HTH and all the best!</p>

DiWarp: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> and <a class="mention" href="/u/dhruvbird">@dhruvbird</a>, for your kind suggestions…</p>
<p>I tried all of them, but still facing the same problem. Here’s what I’ve tried:</p>
<ol>
<li>I changed to a model that was not pre-trained: <code>resnet = models.resnet34(pretrained=False)</code></li>
<li>Changed the device to work on a GPU</li>
<li>Dropped the normalization of (0.5,), (0.5,)</li>
<li>I am applying <code>sigmoid()</code> on the output tensor</li>
<li>I am using now only <strong>one channel</strong> on the images</li>
<li>Changed the num_epochs to 60</li>
<li>I tried the training with 2 different loss functions: <code>L1Loss()</code> and <code>BCELoss()</code></li>
</ol>
<p><strong>With <code>BCELoss()</code>, working on 1 channel, 60 epochs of training:</strong></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/6/36612c1281bf8d9066c0c38aa8b5c81893c25d8a.png" alt="image" data-base62-sha1="7L3YnfJoISPA8nE2b1XDqOPFRcC" width="282" height="217"> <img src="https://discuss.pytorch.org/uploads/default/original/3X/a/8/a81c70eecc5541e70b52be900635c18958a8afbb.png" alt="image" data-base62-sha1="nZb9iexMI4u0REF1FHY6xL8k0qL" width="283" height="217"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/c/9c4911da1ed12419fe82ef0d72de796a06717cb5.png" data-download-href="https://discuss.pytorch.org/uploads/default/9c4911da1ed12419fe82ef0d72de796a06717cb5" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/c/9c4911da1ed12419fe82ef0d72de796a06717cb5.png" alt="image" data-base62-sha1="miz2gAEWb3SeMs2eGXcwdAddZQ1" width="517" height="250" data-dominant-color="A9A9A9"><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">image</span><span class="informations">795×385 14.2 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><strong>With <code>L1Loss()</code>, working on 1 channel, 60 epochs of training:</strong></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/2/62c30aad611f2855745f01a987370c444cc50507.png" alt="image" data-base62-sha1="e5GFm5qA2042TaVmobWEgiuuEvR" width="282" height="217"><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/0/800778269a71793dcd0374e2543eb82849a910b2.png" alt="image" data-base62-sha1="igB7odjakRBtUJHD4vUOLSNwVA6" width="292" height="217"><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/0/50a8e8803b70571699b68721e42a7e5e79f00a47.png" data-download-href="https://discuss.pytorch.org/uploads/default/50a8e8803b70571699b68721e42a7e5e79f00a47" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/0/50a8e8803b70571699b68721e42a7e5e79f00a47.png" alt="image" data-base62-sha1="bvy54ra7uDptcnbXlDSuAwz0DPx" width="517" height="250" data-dominant-color="8A8A8A"><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">image</span><span class="informations">795×385 9.72 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>It’s quite intriguing…</p>

dhruvbird: <aside class="quote no-group" data-username="DiWarp" data-post="1" data-topic="184536">
<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/diwarp/48/52402_2.png" class="avatar"> DiWarp:</div>
<blockquote>
<p><code>transforms.Lambda(lambda x: x + torch.randn_like(x))  # adding noise</code></p>
</blockquote>
</aside>
<p>I’m not sure if this is entirely correct. I mean it adds random noise to your image but changes the range of values from [0.0 … 1.0] to outside this range. Would you consider adding noise that has mean 0.5 and standard deviation 0.5 and then clip the values to [0.0 … 1.0] if they fall outside the range?</p>

dhruvbird: <aside class="quote no-group" data-username="DiWarp" data-post="4" data-topic="184536">
<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/diwarp/48/52402_2.png" class="avatar"> DiWarp:</div>
<blockquote>
<p>I tried the training with 2 different loss functions: <code>L1Loss()</code> and <code>BCELoss()</code></p>
</blockquote>
</aside>
<p>It’s not clear to me how one would use <code>BCELoss()</code> for this problem. Please could you try MSE loss instead?</p>

DiWarp: <p>Hey, <a class="mention" href="/u/dhruvbird">@dhruvbird</a>, take a look. I added this function:</p>
<pre><code class="lang-auto">def add_gaussian_noise(image):
      noise = torch.randn(image.size()) * 0.5 + 0.5
      noisy_image = image + noise
      noisy_image = torch.clamp(noisy_image, 0., 1.)
      return noisy_image
</code></pre>
<p>on:</p>
<pre><code class="lang-auto">noisy_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.float()),  # Convert to FloatTensor
    transforms.Lambda(add_gaussian_noise)  # add noise
])
</code></pre>
<p>and changed the loss function to MSELoss(), here’s the outputs:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/2/c27c8f7adc6174c77ef0ac94a0c5c2c105038ebc.png" alt="image" data-base62-sha1="rKvuVxxqh94ukuupu88XzM0S8Bm" width="287" height="217"><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/1/410efeb9ae4e36b1e1aab0517634eae18d4312c3.png" alt="image" data-base62-sha1="9hxa3JOMmyUz1u96Xp0auqbPmtd" width="292" height="217"><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/6/36aeead764495ef91b7695a1f6a2b438fede1eba.png" data-download-href="https://discuss.pytorch.org/uploads/default/36aeead764495ef91b7695a1f6a2b438fede1eba" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/6/36aeead764495ef91b7695a1f6a2b438fede1eba.png" alt="image" data-base62-sha1="7NKxBOK159Fng8fKfh84DAfHzdM" width="517" height="250" data-dominant-color="B5B5B5"><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">image</span><span class="informations">795×385 13.4 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>