使用AutoDecoder自动解码器实现简单MNIST特征向量提取

AutoDecoder

自动解码器（AD）是论文"DeepSDF: Learning Continuous Signed Distance Functions for Shape Representation" 中使用的一种方法，与传统编码-解码结构不同，AD无编码器，仅有一个解码器。解码器实现特征向量（隐向量）与图片之间的转换。

在训练过程中同时优化特征向量与神经网络参数。如果训练集有N张图片，特征向量长度为n，神经网络参数为m，那么待训练参数共有N*n+m个。

训练完成之后，任给一个特征向量，输入解码器，则可得到一张图片。

DeepSDF原文更为复杂，使用AD生成带符号的最小距离场，进而实现3D形状的生成。DeepSDF原文还有一个针对MNIST手写数据集的简单的案例，但是没有给出源代码，难以上手。笔者在Github上找到了一个具体实现代码https://github.com/alexeybokhovkin/DeepSDF，在此基础上作了一些修改完善，撰写了这篇博客，以作记录，希望对读者有所帮助。

代码

项目共有四个文件，dataset.py用于定义数据集，evaluate.py用于评估训练后的神经网络，network.py用于定义神经网络结构，train.py用于训练神经网络

1 dataset.py

用于导入MNIST数据集（或FashionMNIST）数据集，如果本地没有则会从互联网下载

import torchvision
from torch.utils.data.dataset import Dataset

# A wrapper dataset over MNIST to return images and indices
class DatasetMNIST(Dataset):

    def __init__(self, root_dir, latent_size, transform=None):
        mnist = torchvision.datasets.FashionMNIST(root=root_dir, train=True,download=True)
        self.data = mnist.train_data.float()/255.0

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        image = self.data[index]

        return image.flatten(), index

2 network.py

定义了一个全连接神经网络，输入特征向量，输出图片（展平为向量）。为了便于观察结果，输入的特征向量维数为2。需要注意的是，特征向量也作为被训练参数，共N*n个元素。

import torch
import torch.nn as nn
import torch.nn.init as init

# Autodecoder structure
class AD(nn.Module):
    def __init__(self, image_size=784, z_dim=2, data_shape=60000):
        super(AD, self).__init__()
        self.decoder = nn.Sequential(
            nn.Linear(z_dim, 128),
            nn.ReLU(True),
            nn.Linear(128, 256),
            nn.ReLU(True),
            nn.Linear(256, 512),
            nn.ReLU(True), nn.Linear(512, 28 * 28), nn.Tanh())
        
        self.latent_vectors = nn.Parameter(torch.FloatTensor(data_shape, z_dim))
        
        init.xavier_normal(self.latent_vectors)
    
    def forward(self, ind):
        x = self.latent_vectors[ind]
        return self.decoder(x)
    
    def predict(self, x):
        return self.decoder(x)

3 train.py

使用Adams训练神经网络，batch_size=128，num_epochs=250。训练只需几分钟，完成后保存为model.pth文件，以供调用

import os
import numpy as np
import torch
import torch.nn as nn
from torchvision.utils import save_image
from dataset import DatasetMNIST
from network import AD
from tqdm import tqdm

# Hyper-parameters
image_width = 28
image_size = image_width*image_width
h_dim = 512
num_epochs = 250
batch_size = 128
learning_rate = 1e-3
latent_size = 2

if __name__ == "__main__":
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'

    # Create a directory if not exists
    sample_dir = 'samples'
    os.makedirs(sample_dir, exist_ok=True)

    dataset = DatasetMNIST(root_dir='./data', latent_size=latent_size)

    # Data loader
    data_loader = torch.utils.data.DataLoader(dataset=dataset,
                                            batch_size=batch_size,
                                            shuffle=True)

    model = AD(image_size=image_size, z_dim=latent_size, data_shape=60000).cuda()

    # recusntruction loss
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    model.train()
    for epoch in range(num_epochs):
        tq = tqdm(total=len(data_loader))
        tq.set_description('Epoch {}'.format(epoch))
        for i, (x, ind) in enumerate(data_loader):
            # Forward pass
            x = x.cuda()
            x_reconst = model(ind)

            loss = criterion(x_reconst, x)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            tq.update()
            tq.set_postfix(loss='{:.3f}'.format(loss.item()))

        if epoch%5 == 0:
            with torch.no_grad():
                # Visualize 2D latent space
                steps = 50
                bound = 0.8
                size = image_width
                out = torch.zeros(size=(steps * size, steps * size))

                for i, l1 in enumerate(np.linspace(-bound, bound, steps)):
                    for j, l2 in enumerate(np.linspace(-bound, bound, steps)):
                        vector = torch.tensor([l1, l2]).to(dtype=torch.float32).cuda()
                        out_ = model.predict(vector)
                        out[i * size:(i + 1) * size, j * size:(j + 1)
                            * size] = out_.view(size, size)
                save_image(out, os.path.join(
                    sample_dir, 'latent_space-{}.png'.format(epoch + 1)))


    # save model
    torch.save(model, 'model.pth')

4 evaluate.py

加载训练好的模型，遍历特征向量[l1,l2]，使用解码器生成对应的图像，保存为latent_space-eval.png

将所有训练集的特征向量[l1,l2]绘制在一张图上，保存为latent_space-distribution.png

最后将两张图片拼接得到latent_space-merged.png

import os
import torch
import numpy as np
from torchvision.utils import save_image
from train import image_width

if __name__ == "__main__":
    # Create a directory if not exists
    sample_dir = 'samples'
    os.makedirs(sample_dir, exist_ok=True)

    # 选择GPU或CPU
    device = torch.device(
        "cuda") if torch.cuda.is_available() else torch.device("cpu")

    # 从文件加载已经训练完成的模型
    model = torch.load('model.pth', map_location=device)
    model.eval()  # 设置模型为evaluation状态

    print(model)

    # Visualize 2D latent space
    steps = 50
    bound = 0.8
    size = image_width
    out_grid = torch.zeros(size=(steps * size, steps * size))

    for i, l1 in enumerate(np.linspace(-bound, bound, steps)):
        for j, l2 in enumerate(np.linspace(-bound, bound, steps)):
            vector = torch.tensor([l1, l2]).to(dtype=torch.float32).cuda()
            out_ = model.predict(vector)
            out_grid[i * size:(i + 1) * size, j * size:(j + 1)
                * size] = out_.view(size, size)
    save_image(out_grid, os.path.join(
        sample_dir, 'latent_space-eval.png'))
    
    out_dist = torch.zeros(size=(steps * size, steps * size))

    latent_vectors_scaled=model.latent_vectors.cpu().detach().numpy()
    latent_vectors_scaled=np.clip(latent_vectors_scaled, -bound+0.005, bound-0.005)
    latent_vectors_scaled = ((latent_vectors_scaled+bound)/(2.0*bound)*steps*size*1.0)

    for i in range(len(latent_vectors_scaled)):
        l1=round(latent_vectors_scaled[i][0])
        l2=round(latent_vectors_scaled[i][1])
        out_dist[l1, l2]=1.0
        out_dist[l1-1, l2]=1.0
        out_dist[l1+1, l2]=1.0
        out_dist[l1, l2-1]=1.0
        out_dist[l1-1, l2-1]=1.0
        out_dist[l1+1, l2-1]=1.0
        out_dist[l1, l2+1]=1.0
        out_dist[l1-1, l2+1]=1.0
        out_dist[l1+1, l2+1]=1.0

    save_image(out_dist, os.path.join(
    sample_dir, 'latent_space-distribution.png'))

    out_merged = torch.cat((out_grid, out_dist), dim=1)
    save_image(out_merged, os.path.join(
    sample_dir, 'latent_space-merged.png'))

    print(model.latent_vectors.max(), model.latent_vectors.min())

结果

训练前遍历特征向量绘制得到解码后的图片:

epoch=16, 遍历特征向量绘制得到解码后的图片:

epoch=101, 遍历特征向量绘制得到解码后的图片:
epoch=250, 遍历特征向量绘制得到解码后的图片:
训练样本分布如下：

结论

（1） AD事实上是一种特征提取方法，本文从数据集中提取了一个2D特征，在2D平面内重构出了原始数据集。在实际使用中，特征向量是一个高维向量，效果会更好。
（2） 特征向量的分布近似于一个正态分布，但是不同类别之间存在鸿沟
（3）AD的本质是对图片数据进行压缩，图片公有信息蕴含于神经网络参数中，个体信息蕴含于特征向量。