def test(Data, model, n = 6):
"""
Data:'MNIST' or 'CIFAR'
"""
if Data == 'MNIST':
_, test_loader = DATA.MNISTLoader(p, download = False, batch_size=batch_size)
elif Data == 'CIFAR':
_, test_loader = DATA.CIFAR10Loader(p, download = False, batch_size=batch_size)
else: raise ValueError("Data set not support")
model = model.eval()
with torch.no_grad():
idx, (X,_) = next(enumerate(test_loader))
X = X.to(device)
X_bar, mu, logvar = model.forward(X)
if Data == 'MNIST':
img = torch.cat([X[:n], X_bar.view(batch_size, 1, 28, 28)[:n]])
plt.pause(0.001)
plt.imshow(img.view(28*n*2,28*1).cpu())
#save_image(img.cpu(),
#'C:/Users/funrr/Desktop/DL3/results/test0' + '.png', nrow=n)
elif Data == 'CIFAR':
img = torch.cat([X[:n], X_bar.view(batch_size, 3, 32, 32)[:n]])
plt.pause(0.001)
plt.imshow(np.moveaxis(img.cpu().numpy(), 1, 3).reshape(n*2*32,32,3))
import os
import torch
import torch.nn as nn
import torch.optim as optim
import DATA
import VAE
import utensils
import matplotlib.pyplot as plt
os.chdir(r"C:\Users\funrr\Desktop\DL3")
work_dir = os.getcwd()
p = os.path.join(work_dir, 'DATA')
batch_size = 300
train_loader, _ = DATA.MNISTLoader(p, download=False, batch_size=batch_size)
#idx, (X,_) = next(enumerate(train_loader))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
vae = VAE.MLP_VAE(Data='MNIST', batch_size=batch_size, z_dim=30).to(device)
optimizer = optim.Adam(vae.parameters(),
lr=0.01,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0,
amsgrad=False)
loss_log = []
for epoch in range(20):
for step, (train_data, _) in enumerate(train_loader):