1. AutoEncoders回顾

如下图：AutoEncoders实际上就是重建自己的过程

2. 实现网络结构

创建类继承自nn.Model，并实现init和forward函数，init中实现encoder、decoder
直接上代码，其中encoder维度变化784->256->64->20，decoder是相反过程

import  torch
from    torch import nnclass AE(nn.Module):def __init__(self):super(AE, self).__init__()# [b, 784] => [b, 20]self.encoder = nn.Sequential(nn.Linear(784, 256),nn.ReLU(),nn.Linear(256, 64),nn.ReLU(),nn.Linear(64, 20),nn.ReLU())# [b, 20] => [b, 784]self.decoder = nn.Sequential(nn.Linear(20, 64),nn.ReLU(),nn.Linear(64, 256),nn.ReLU(),nn.Linear(256, 784),nn.Sigmoid())def forward(self, x):""":param x: [b, 1, 28, 28]:return:"""batchsz = x.size(0)# flattenx = x.view(batchsz, 784)# encoderx = self.encoder(x)# decoderx = self.decoder(x)# reshapex = x.view(batchsz, 1, 28, 28)return x, None

3. 实现main函数

• 加载minist数据集
• 初始化网络
• 初始化loss函数
• 初始化优化器
• 循环迭代数据输入模型、计算loss、优化器优化参数
  完整代码

import  torch
from    torch.utils.data import DataLoader
from    torch import nn, optim
from    torchvision import transforms, datasetsfrom    ae import AE
#from    vae import VAEimport  visdomdef main():mnist_train = datasets.MNIST('mnist', True, transform=transforms.Compose([transforms.ToTensor()]), download=True)mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True)mnist_test = datasets.MNIST('mnist', False, transform=transforms.Compose([transforms.ToTensor()]), download=True)mnist_test = DataLoader(mnist_test, batch_size=32, shuffle=True)x, _ = iter(mnist_train).next()print('x:', x.shape)device = torch.device('cuda')model = AE().to(device)#model = VAE().to(device)criteon = nn.MSELoss()optimizer = optim.Adam(model.parameters(), lr=1e-3)print(model)viz = visdom.Visdom()for epoch in range(1000):for batchidx, (x, _) in enumerate(mnist_train):# [b, 1, 28, 28]x = x.to(device)x_hat, _ = model(x)loss = criteon(x_hat, x)# backpropoptimizer.zero_grad()loss.backward()optimizer.step()print(epoch, 'loss:', loss.item(), 'kld:', kld.item())x, _ = iter(mnist_test).next()x = x.to(device)with torch.no_grad():x_hat, kld = model(x)viz.images(x, nrow=8, win='x', opts=dict(title='x'))viz.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat'))if __name__ == '__main__':main()