def main():
G = Generator(z_dim=20)
D = Discriminator(z_dim=20)
E = Encoder(z_dim=20)
G.apply(weights_init)
D.apply(weights_init)
E.apply(weights_init)
train_img_list=make_datapath_list(num=200)
mean = (0.5,)
std = (0.5,)
train_dataset = GAN_Img_Dataset(file_list=train_img_list, transform=ImageTransform(mean, std))
batch_size = 64
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
num_epochs = 1500
G_update, D_update, E_update = train_model(G, D, E, dataloader=train_dataloader, num_epochs=num_epochs, save_model_name='Efficient_GAN')
if device == 'cuda':
G = torch.nn.DataParallel(G)
D = torch.nn.DataParallel(D)
print("parallel mode")
batch_size = 8
z_dim = 20
fixed_z = torch.randn(batch_size, z_dim)
fixed_z = fixed_z.view(fixed_z.size(0), fixed_z.size(1), 1, 1)
fake_images = G(fixed_z.to(device))
train_img_list = make_datapath_list(num=1000)
mean = (0.5, )
std = (0.5, )
train_dataset = GAN_Img_Dataset(file_list=train_img_list,
transform=ImageTransform(mean, std))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
batch_iterator = iter(train_dataloader)
# fetch first element
images = next(batch_iterator)
fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
#train is upside
plt.subplot(2, 5, i + 1)
plt.imshow(images[i][0].cpu().detach().numpy(), 'gray')