Pytorch实现VAE变分自动编码器生成MNIST手写数字图像

import os import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.autograd import Variable from torchvision import transforms import torchvision.datasets as dst from torchvision.utils import save_image from model import VAE # 训练超参数配置 epochs = 50 batch_size = 64 num_workers = 0 log_interval = 10 latent_dim = 32 # 损失函数 def loss_func(recon_x, x, mu, logvar): BCE = F.binary_cross_entropy(recon_x, x, size_average=False) KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) return BCE + KLD if __name__ == '__main__': # 加载数据集 transform = transforms.Compose([transforms.ToTensor()]) data_train = dst.MNIST('MNIST_data/', train=True, transform=transform, download=True) data_test = dst.MNIST('MNIST_data/', train=False, transform=transform, download=True) train_loader = torch.utils.data.DataLoader(dataset=data_train, num_workers=num_workers, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=data_test, num_workers=num_workers, batch_size=batch_size, shuffle=True) # 创建VAE模型 vae = VAE(latent_dim).cuda() # 创建优化器 optimizer = optim.Adam(vae.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) os.makedirs('result', exist_ok=True) # 创建文件夹 # 开始训练 for epoch in range(1, epochs): vae.train() total_loss = 0 for i, (data, _) in enumerate(train_loader, 0): data = Variable(data).cuda() optimizer.zero_grad() recon_x, mu, logvar = vae.forward(data) loss = loss_func(recon_x, data, mu, logvar) loss.backward() total_loss += loss.item() optimizer.step() if i % log_interval == 0: sample = Variable(torch.randn(64, latent_dim)).cuda() sample = vae.decoder(vae.fc2(sample).view(64, 128, 7, 7)).cpu() save_image(sample.data.view(64, 1, 28, 28), 'result/sample_' + str(epoch) + '.png') print('Train Epoch:{} -- [{}/{} ({:.0f}%)] -- Loss:{:.6f}'.format( epoch, i*len(data), len(train_loader.dataset), 100.*i/len(train_loader), loss.item()/len(data))) print('====> Epoch: {} Average loss: {:.4f}'.format( epoch, total_loss / len(train_loader.dataset))) torch.save(vae.state_dict(), 'model.pth')