<jupyter_start><jupyter_text><jupyter_code>import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as f

import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
data_transform = transforms.ToTensor()
dataset = torchvision.datasets.MNIST('./root',download=True,transform = data_transform,train=True)
testset = torchvision.datasets.MNIST('./root',download=True,transform = data_transform,train=False)
batch_size = 20
num_workers = 0
trainloader = DataLoader(dataset , shuffle = True , batch_size = batch_size , num_workers = num_workers)
testloader  = DataLoader(testset , shuffle = True , batch_size = batch_size , num_workers = num_workers)
images , labels = iter(trainloader).next()
images = images.numpy()

fig = plt.figure(figsize = (25,4))

for i in range(batch_size):
  ax = fig.add_subplot(2 , batch_size / 2 , i+1 , xticks = [] , yticks = [])
  ax.imshow(np.squeeze(images[i]) , cmap = 'gray')
  ax.set_title(labels[i].item())
images , labels = iter(trainloader).next()
images = images.numpy()

noise_factor = 0.5

fig = plt.figure(figsize = (25,4))

for i in range(batch_size):
  images_n = images[i] + np.random.randn(*images[i].shape) * noise_factor
  ax = fig.add_subplot(2 , batch_size / 2 , i+1 , xticks = [] , yticks = [])
  ax.imshow(np.squeeze(images_n) , cmap = 'gray')
  ax.set_title(labels[i].item())

images.shape
import torch.nn as nn
import torch.nn.functional as F

# define the NN architecture
class ConvDenoiser(nn.Module):
    def __init__(self):
        super(ConvDenoiser, self).__init__()
        ## encoder layers ##
        # conv layer (depth from 1 --> 32), 3x3 kernels
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)  
        # conv layer (depth from 32 --> 16), 3x3 kernels
        self.conv2 = nn.Conv2d(32, 16, 3, padding=1)
        # conv layer (depth from 16 --> 8), 3x3 kernels
        self.conv3 = nn.Conv2d(16, 8, 3, padding=1)
        # pooling layer to reduce x-y dims by two; kernel and stride of 2
        self.pool = nn.MaxPool2d(2, 2)
        
        ## decoder layers ##
        # transpose layer, a kernel of 2 and a stride of 2 will increase the spatial dims by 2
        self.t_conv1 = nn.ConvTranspose2d(8, 8, 3, stride=2)  # kernel_size=3 to get to a 7x7 image output
        # two more transpose layers with a kernel of 2
        self.t_conv2 = nn.ConvTranspose2d(8, 16, 2, stride=2)
        self.t_conv3 = nn.ConvTranspose2d(16, 32, 2, stride=2)
        # one, final, normal conv layer to decrease the depth
        self.conv_out = nn.Conv2d(32, 1, 3, padding=1)


    def forward(self, x):
        ## encode ##
        # add hidden layers with relu activation function
        # and maxpooling after
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        # add second hidden layer
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        # add third hidden layer
        x = F.relu(self.conv3(x))
        x = self.pool(x)  # compressed representation
        
        ## decode ##
        x = F.relu(self.t_conv1(x))
        x = F.relu(self.t_conv2(x))
        x = F.relu(self.t_conv3(x))
        # transpose again, output should have a sigmoid applied
        x = F.sigmoid(self.conv_out(x))
                
        return x

# initialize the NN
model = ConvDenoiser()
model.cuda()
print(model)
# specify loss function
criterion = nn.MSELoss().cuda()

# specify loss function
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
noise_factor = 0.5
n_epochs = 10
trainloss = []

for epoch in range(n_epochs):
  train_loss = 0.0
  for batch_i , (images , target) in enumerate(trainloader):
    noisy_images = images + torch.randn(*images.shape) * noise_factor
    noisy_images = torch.tensor(noisy_images).cuda()
    optimizer.zero_grad()
    output = model(noisy_images)
    loss = criterion(output , images.cuda())
    loss.backward()
    optimizer.step()
    train_loss += loss.item()
    
  print("Epoch: {} , Batch: {} , Training Loss: {}".format(epoch + 1 , batch_i + 1 , train_loss / len(trainloader)))
  trainloss.append(train_loss / len(trainloader))

print("Training Completed !!!")
images , labels = iter(testloader).next()
images = images.numpy()
for i in range(batch_size):
  images_n = images + np.random.randn(*images.shape) * noise_factor
  preds = model(torch.cuda.FloatTensor(images_n).cuda()).cpu().detach().numpy()

noise_factor = 0.5

fig = plt.figure(figsize = (25,4))

for i in range(batch_size):
  images_n = images + np.random.randn(*images.shape) * noise_factor
  ax = fig.add_subplot(2 , batch_size / 2 , i+1 , xticks = [] , yticks = [])
  ax.imshow(preds[i].reshape(28,28) , cmap = 'gray')<jupyter_output>/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:1709: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.
  warnings.warn("nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.")
StarcoderdataJupyterStructuredCleanDedup