VAE模型以及利用MINIST训练生成数字（效果较好）

from torch import nn import torch class VAE(nn.Module): def __init__(self, image_size=28, z_dim=100) -> None: super().__init__() self.encoder = nn.Sequential( # 由若干卷积和一个全连接构成encoder nn.Conv2d(in_channels=1, out_channels=20, kernel_size=3), nn.ReLU(), nn.Conv2d(in_channels=20, out_channels=20, kernel_size=3), nn.ReLU(), nn.Flatten(start_dim=1), nn.Linear(in_features=11520, out_features=800), nn.ReLU() ) self.sigma = nn.Linear(800, z_dim) self.mu = nn.Linear(800, z_dim) self.sigmoid = nn.Sigmoid() self.decoder = nn.Sequential( # 两个全连接构成decoder nn.Linear(in_features=z_dim, out_features=2000), nn.ReLU(), nn.Linear(in_features=2000, out_features=image_size * image_size) ) def reparameterize(self, sigma, mu): # 冲参数计算，计算z sigma = torch.exp(sigma / 2) # 这里sigma可能非正，所以加了exp，然后开根号 eps = torch.rand_like(sigma) # 抽样得到epsilon return mu + eps * sigma def encode(self, x): hidden = self.encoder(x) sigma, mu = self.sigma(hidden), self.mu(hidden) return sigma, mu def decode(self, z): return self.decoder(z) def forward(self, x): sigma, mu = self.encode(x) z = self.reparameterize(sigma, mu) return sigma, mu, self.sigmoid(self.decode(z)) # 这里的sigmoid是为了计算交叉熵损失 def loss_function(sigma, mu, x, reconstruct): BCE_loss = nn.BCELoss(reduction='sum') kl_divergece = -0.5 * torch.sum(1 + sigma - torch.exp(sigma) - mu ** 2) # 这里的sigma本来应该是正的，所以用exp来代替了 x = torch.flatten(x, start_dim=1) cross = BCE_loss(reconstruct, x) # 交叉熵部分 return kl_divergece + cross import torch import torchvision from torch.utils.data import DataLoader from torchvision.utils import save_image if __name__ == '__main__': BATCH_SIZE = 150 LR = 0.001 EPOCH = 100 model = VAE().cuda() optim = torch.optim.Adam(lr=LR, params=model.parameters()) data = torchvision.datasets.MNIST('./data', train=True, transform=torchvision.transforms.ToTensor(), download=True) dataloader = DataLoader(dataset=data, batch_size=BATCH_SIZE, shuffle=True) for epoch in range(EPOCH): for i, (imgs, labels) in enumerate(dataloader): imgs = imgs.cuda() sigma, mu, recon = model(imgs) loss = loss_function(sigma, mu, imgs, recon) optim.zero_grad() loss.backward() optim.step() if i % 50 == 0: print('loss', loss.item()) fake_images = recon.view(-1, 1, 28, 28)[:5] x_concat = torch.cat([imgs.view(-1, 1, 28, 28)[:5, :], fake_images], dim=3) save_image(x_concat, 'generate_images/generate_images-{}-{}.png'.format(epoch + 1, i)) torch.save(model.state_dict(), './model.pkl')