1.Auto Encoder
具体流程:输入==RESHAPE==>784=>1000=>1000=>20=>1000=>1000=>784==RESHAPE==>输出
1.网络
网络层:encoder{[b, 784] => [b, 20]} + decoder{[b, 20] => [b, 784]}
连接层:input.reshape->encoder->decoder->output.reshpe
class 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() #输出压缩到0~1
)
def forward(self, x):
"""param x: [b, 1, 28, 28]"""
batchsz = x.size(0)
# flatten
x = x.view(batchsz, 784)
# encoder
x = self.encoder(x)
# decoder
x = self.decoder(x)
# reshape
x = x.view(batchsz, 1, 28, 28)
return x, None
2.训练&测试:
def main():
mnist_test = DataLoader(mnist_test, batch_size=32, shuffle=True)
x, _ = iter(mnist_train).next()
model = VAE().to(device)
criteon = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
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)
# backprop
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(epoch, 'loss:', loss.item())
2.Variational Autoencoder 变分自动编码器
1.网络
网络层:encoder{[b, 784] => [b, 20]} + [b,20]=>μ[b,10]+σ[b,10] +decoder{[b, 10] => [b, 784]}
连接层:input.reshape->encoder->->decoder->计算KL->output.reshpe
class VAE(nn.Module):
def __init__(self):
super(VAE, self).__init__()
# [b, 784] => [b, 10]
# sigma: [b, 10]
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(10, 64),
''''''
nn.ReLU(),
nn.Linear(64, 256),
nn.ReLU(),
nn.Linear(256, 784),
nn.Sigmoid()
)
self.criteon = nn.MSELoss()
def forward(self, x):
"""param x: [b, 1, 28, 28]"""
batchsz = x.size(0)
# flatten
x = x.view(batchsz, 784)
'''修改部分'''
# encoder
# [b, 20], 包含mean和σ
h_ = self.encoder(x)
# [b, 20] => [b, 10] and [b, 10]
mu, sigma = h_.chunk(2, dim=1)
# reparametrize trick, epison~N(0, 1)
h = mu + sigma * torch.randn_like(sigma)
''''''
# decoder
x_hat = self.decoder(h)
# reshape
x_hat = x_hat.view(batchsz, 1, 28, 28)
'''KL'''
kld = 0.5 * torch.sum(
torch.pow(mu, 2) +
torch.pow(sigma, 2) -
torch.log(1e-8 + torch.pow(sigma, 2)) - 1
) / (batchsz*28*28)
''''''
return x_hat, kld
2.训练&测试:
注意loss计算:
for epoch in range(1000):
for batchidx, (x, _) in enumerate(mnist_train):
# [b, 1, 28, 28]
x = x.to(device)
x_hat, kld = model(x)
''''''
loss = criteon(x_hat, x)
elbo = - loss - 1.0 * kld
loss = - elbo
''''''
# backprop
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(epoch, 'kld loss:', kld.item())