实战PyTorch（三）：Auto Encoder & Variational Autoencoder_variational autoencoder pytorch

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())