AutoEncoder是一种无监督学习方式,包含编码器和解码器两部分组成。AutoEncoder常被用于降维,关于降维可参考该网站。
AutoEncoder本质上还是一种线性变换,实现起来也比较简单。本文使用AutoEncoder来对MNIST数据集进行学习和生成。
代码如下:
首先是网络模型的定义,命名为auto_encoder.py
import torch
from torch import nn
class AutoEncoder(nn.Module):
def __init__(self):
super(AutoEncoder, 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):
# x: [b, 1, 28, 28]
batch_size = x.size(0)
# flatten
x = x.view(batch_size, -1)
# encoder
x = self.encoder(x)
# decoder
x = self.decoder(x)
# reshape
x = x.view(batch_size, 1, 28, 28)
return x
接下来写主函数main().py
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from auto_encoder import AutoEncoder # 从上一个程序中导入网络模型
import visdom # 使用visdom做可视化
def main():
# 加载数据集
mnist_train = datasets.MNIST(root='mnist', train=True, transform=transforms.Compose([
transforms.ToTensor()]), download=True)
mnist_train_loader = DataLoader(mnist_train, batch_size=32, shuffle=True)
mnist_test = datasets.MNIST(root='mnist', train=True, transform=transforms.Compose([
transforms.ToTensor()]), download=True)
mnist_test_loader = DataLoader(mnist_test, batch_size=32, shuffle=False)
x, _ = iter(mnist_train_loader).next() # 因为是无监督学习 所以不使用label标签
print("x.shape: ", x.shape)
# 定义损失 选择优化器
device = torch.device('cuda')
model = AutoEncoder().to(device) # 不要忘了转移到GPU上
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
print(model)
viz = visdom.Visdom() # 实例化visdom
# train
for epoch in range(50):
for batch_idx, (x, _) in enumerate(mnist_train_loader):
# x: [b, 1, 28, 28]
x = x.to(device)
x_hat = model(x)
loss = criterion(x_hat, x)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Epoch: ", epoch, "Loss: ", loss.item())
x, _ = iter(mnist_test_loader).next() # 测试集上验证
x = x.to(device)
with torch.no_grad():
x_hat = 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'))
# 每行画8个图
if __name__ == '__main__':
main()