Dimensionality Reduction and Scattered Data Visualization with MNIST

We live in 3D world, and we can only view scattered data in 1D, 2D, or 3D. Yet, we deal with data that have very large dimension. Consider MNIST dataset, which is considered to be a toy example in deep learning field, consists of 28 X 28 gray images; that is 784 dimensions.

How would this MNIST data look like in 2D or 3D after dimensionality reduction? Let's figure it out! I am going to write the code in Pytorch. I have to say, Pytorch is so much better than other deep learning libraries, such as Theano or Tensorflow. Of course, it is just my personal opinion, so let's not get into this argument here.

What I want to do is to take Pytorch's MNIST example found here, and make some modifications to reduce the data dimension to 2D and plot scattered data. This will be a very good example that shows how to do all the following in Pytorch:
1. Create a custom network
2. Create a custom layer
3. Transfer learning from an existing model
4. Save and load a model

Here is the plot I get from running the code below.

This code is tested on Pytorch 0.3.1.