This work is based on the project of TommyZihao - zihaopytorch
The main purpose of this work is to convert his code to run on the GPU, and this work also converts some training, testing and visualization processes into the form of Python functions for easy invocation.