概念
无监督学习,用于压缩还原、在保留部分原始信息的基础上生成、换脸……
小明是个画家,有 A 和 B 两个人,看 A 两秒钟,看 B 一小时,然后画 A,画中会有 B 的特征。
编码:提取特征。(看 A)
解码:学习新特征,还原。(看 B 画 A)
编码结果向 N(0, 1) 靠近,因为真实分布和 N(0, 1) 是有差别的,所以解码生成的图片会模糊。
编码器得出的正态分布:N(μ, σ²)。
标准正态分布:N(0, 1)。
关系:N(μ, σ²) = N(0, 1) * σ² + μ。
编码损失:计算 N(μ, σ²) 和 N(0, 1) 差距,也就是 KL 散度。
KL(N(μ, σ²) || N(0, 1)) = (-logσ² + μ² + σ² - 1) / 2。
实验(手写数字图片压缩还原)
数据集:MNIST。
网络结构:
- 编码:RNN + 标准化(BN)+ 激活(ReLU)。
- 解码:转置卷积 + 标准化(BN)+ 激活(ReLU)。
优化器:Adam。
损失函数:均方差(MSELoss)。
输出:图片。
网络
import torch
from torch import nn
class Encoder(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Sequential(
# (28 * 28)→(14 * 14)
nn.Conv2d(1, 128, 3, 2, 1), nn.BatchNorm2d(128), nn.ReLU(),
# (14 * 14)→(7 * 7)
nn.Conv2d(128, 256, 3, 2, 1), nn.BatchNorm2d(256), nn.ReLU(),
# (7 * 7)→(3 * 3)
nn.Conv2d(256, 512, 3, 2, 0), nn.BatchNorm2d(512), nn.ReLU(),
# (3 * 3)→(1 * 1)
nn.Conv2d(512, 2, 3, 1, 0)
)
def forward(self, x):
# [n,2,1,1]
out = self.conv(x)
# 取 c 轴第一个作为 μ
miu = out[:, :1, :, :]
# 取 c 轴第二个作为 logσ²
log_sigma = out[:, 1:, :, :]
return miu, log_sigma
class Decoder(nn.Module):
def __init__(self):
super().__init__()
self.conv_t = nn.Sequential(
nn.ConvTranspose2d(128, 512, 3, 1, 0), nn.BatchNorm2d(512), nn.ReLU(),
nn.ConvTranspose2d(512, 256, 3, 2, 0), nn.BatchNorm2d(256), nn.ReLU(),
# 尺寸除不尽时,要加 output_padding = 1
nn.ConvTranspose2d(256, 128, 3, 2, 1, 1), nn.BatchNorm2d(128), nn.ReLU(),
nn.ConvTranspose2d(128, 1, 3, 2, 1, 1)
)
def forward(self, miu, log_sigma, z):
# σ² = e ** logσ²
x = z * torch.exp(log_sigma) + miu
# [n,h,w,c] → [n,c,h,w]
x = x.permute(0, 3, 1, 2)
return self.conv_t(x)
class MyNet(nn.Module):
def __init__(self):
super().__init__()
self.encoder = Encoder()
self.decoder = Decoder()
def forward(self, x, z):
miu, log_sigma = self.encoder(x)
out = self.decoder(miu, log_sigma, z)
return miu, log_sigma, out
训练
import torch
from torch import nn
from torchvision import datasets, transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
import os
from net import MyNet
batch_size = 100
net_path = r"modules/mynet.pth"
train_flag = True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
if train_flag:
dataset = datasets.MNIST(r"data", train=True, transform=transform, download=True)
dataloader = DataLoader(dataset, batch_size, shuffle=True)
else:
dataset = datasets.MNIST(r"data", train=False, transform=transform, download=False)
dataloader = DataLoader(dataset, batch_size, shuffle=False)
if __name__ == '__main__':
# 加载网络
if os.path.isfile(net_path):
net = torch.load(net_path).to(device)
else:
net = MyNet().to(device)
opt = torch.optim.Adam(net.parameters())
# 只在 bitch 上做平均
loss_fn = nn.MSELoss(reduction='sum')
if train_flag:
# 训练
net.train()
while True:
for i, (x, y) in enumerate(dataloader):
x = x.to(device)
# 从标准正态分布中取值
z = torch.randn(128).to(device)
miu, log_sigma, out = net(x, z)
# 编码损失(N(μ,σ²) 和 N(0,1))
loss_e = torch.mean((-log_sigma + torch.exp(log_sigma) + miu ** 2 - 1) / 2)
# 解码损失(生成的图和原图)
loss_d = loss_fn(out, x)
loss = loss_e + loss_d
opt.zero_grad()
loss.backward()
opt.step()
print("i:{},loss:{:.5}".format(i, loss))
# 保存网络
torch.save(net, net_path)
else:
# 测试
net.eval()
for i, (x, y) in enumerate(dataloader):
x = x.to(device)
# 从标准正态分布中取值
z = torch.randn(128).to(device)
miu, log_sigma, out = net(x, z)
# 保存图片
fake_img = out.data
img = x.data
save_image(fake_img, "./img/{}-fake_img.png".format(i), nrow=10)
save_image(img, "./img/{}-real_img.png".format(i), nrow=10)