Denoising Autoencoder

graph = Graph(format="png")

# Input layer
graph.node("a", "28x28x1 Input")

# the Encoder
graph.node("b", "28x28x32 Convolution")
graph.node("c", "14x14x32 MaxPool")
graph.node("d", "14x14x16 Convolution")
graph.node("e", "7x7x16 MaxPool")
graph.node("f", "7x7x8 Convolution")
graph.node("g", "3x3x8 MaxPool")

# The Decoder
graph.node("h", "7x7x8 Transpose Convolution")
graph.node("i", "14x14x16 Transpose Convolution")
graph.node("j", "28x28x32 Transpose Convolution")
graph.node("k", "28x28x1 Convolution")

# The Output
graph.node("l", "28x28x1 Output")

edges = "abcdefghijkl"
graph.edges([edges[edge] + edges[edge+1] for edge in range(len(edges) - 1)])

graph.render("graphs/network.dot")
graph

Layer = namedtuple("Layer", "kernel stride in_depth out_depth padding".split())
Layer.__new__.__defaults__= (0,)
def output_size(input_size: int, layer: Layer, expected: int) -> int:
    """Calculates the output size of the layer

    Args:
     input_size: the size of the input to the layer
     layer: named tuple with values for the layer
     expected: the value you are expecting

    Returns:
     the size of the output

    Raises:
     AssertionError: the calculated value wasn't the expected one
    """
    size = 1 + int(
        (input_size - layer.kernel + 2 * layer.padding)/layer.stride)
    print(layer)
    print("Layer Output: {0} x {0} x {1}".format(size, layer.out_depth))
    assert size == expected, size
    return size

model = ConvDenoiser()
print(model)

test = ConvDenoiser()
dataiter = iter(train_loader)
images, labels = dataiter.next()
x = test.convolution_1(images)
assert x.shape == torch.Size([BATCH_SIZE, 32, 28, 28])
print(x.shape)

x = test.max_pool(x)
assert x.shape == torch.Size([BATCH_SIZE, 32, 14, 14])
print(x.shape)

x = test.convolution_2(x)
assert x.shape == torch.Size([BATCH_SIZE, 16, 14, 14])
print(x.shape)

x = test.max_pool(x)
assert x.shape == torch.Size([BATCH_SIZE, 16, 7, 7])
print(x.shape)

x = test.convolution_3(x)
assert x.shape == torch.Size([BATCH_SIZE, 8, 7, 7])
print(x.shape)

x = test.max_pool(x)
assert x.shape == torch.Size([BATCH_SIZE, 8, 3, 3]), x.shape

x = test.transpose_convolution_1(x)
assert x.shape == torch.Size([BATCH_SIZE, 8, 7, 7]), x.shape
print(x.shape)

x = test.transpose_convolution_2(x)
assert x.shape == torch.Size([BATCH_SIZE, 16, 14, 14])
print(x.shape)

x = test.transpose_convolution_3(x)
assert x.shape == torch.Size([BATCH_SIZE, 32, 28, 28])
print(x.shape)

x = test.convolution_out(x)
assert x.shape == torch.Size([BATCH_SIZE, 1, 28, 28])
print(x.shape)

loss = criterion(outputs, images)

class Trainer:
    """Trains our model

    Args:
     data: data-iterator for training
     epochs: number of times to train on the data
     noise: factor for the amount of noise to add
     learning_rate: rate for the optimizer
    """
    def __init__(self, data: torch.utils.data.DataLoader, epochs: int=30,
                 noise:float=0.5,
                 learning_rate:float=0.001) -> None:
        self.data = data
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.noise = noise
        self._criterion = None
        self._model = None
        self._device = None
        self._optimizer = None
        return

    @property
    def device(self) -> torch.device:
        """CUDA or CPU"""
        if self._device is None:
            self._device = torch.device(
                "cuda:0" if torch.cuda.is_available() else "cpu")
        return self._device

    @property
    def criterion(self) -> nn.MSELoss:
        """Loss-calculator"""
        if self._criterion is None:
            self._criterion = nn.MSELoss()
        return self._criterion

    @property
    def model(self) -> ConvDenoiser:
        """Our model"""
        if self._model is None:
            self._model = ConvDenoiser()
            self.model.to(self.device)
        return self._model

    @property
    def optimizer(self) -> torch.optim.Adam:
        """The gradient descent optimizer"""
        if self._optimizer is None:
            self._optimizer = torch.optim.Adam(self.model.parameters(),
                                               lr=self.learning_rate)
        return self._optimizer

    def __call__(self) -> None:
        """Trains the model on the data"""
        self.model.train()
        started = datetime.now()
        for epoch in range(1, self.epochs + 1):
            train_loss = 0.0
            for batch in self.data:
                images, _ = batch
                images = images.to(self.device)
                ## add random noise to the input images
                noisy_imgs = (images
                              + self.noise
                              * torch.randn(*images.shape).to(self.device))
                # Clip the images to be between 0 and 1
                noisy_imgs = numpy.clip(noisy_imgs, 0., 1.).to(self.device)

                # clear the gradients of all optimized variables
                self.optimizer.zero_grad()
                ## forward pass: compute predicted outputs by passing *noisy* images to the model
                outputs = self.model(noisy_imgs)
                # calculate the loss
                # the "target" is still the original, not-noisy images
                loss = self.criterion(outputs, images)
                # backward pass: compute gradient of the loss with respect to model parameters
                loss.backward()
                # perform a single optimization step (parameter update)
                self.optimizer.step()
                # update running training loss
                train_loss += loss.item() * images.size(0)

            # print avg training statistics 
            train_loss = train_loss/len(train_loader)
            print('Epoch: {} \tTraining Loss: {:.6f}'.format(
                epoch, 
                train_loss
                ))
        ended = datetime.now()
        print("Ended: {}".format(ended))
        print("Elapsed: {}".format(ended - started))
        return

train_the_model = Trainer(train_loader)
train_the_model()

# obtain one batch of test images
dataiter = iter(test_loader)
images, labels = dataiter.next()

# add noise to the test images
noisy_imgs = images + noise_factor * torch.randn(*images.shape)
noisy_imgs = numpy.clip(noisy_imgs, 0., 1.)

# get sample outputs
noisy_imgs = noisy_imgs.to(train_the_model.device)
output = train_the_model.model(noisy_imgs)
# prep images for display
noisy_imgs = noisy_imgs.cpu().numpy()

# output is resized into a batch of iages
output = output.view(BATCH_SIZE, 1, 28, 28)
# use detach when it's an output that requires_grad
output = output.detach().cpu().numpy()

# plot the first ten input images and then reconstructed images
fig, axes = pyplot.subplots(nrows=2, ncols=10, sharex=True, sharey=True, figsize=(25,4))

# input images on top row, reconstructions on bottom
for noisy_imgs, row in zip([noisy_imgs, output], axes):
    for img, ax in zip(noisy_imgs, row):
        ax.imshow(numpy.squeeze(img), cmap='gray')
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)