GAN|生成人脸|全连接网络|代码

import torchvision.datasets
mnist_dataset = torchvision.datasets.CelebA(root='.', download=True)

import h5py
import numpy
import matplotlib.pyplot as plt
import zipfile
import imageio
import os
import torch
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import pandas
if torch.cuda.is_available():
  torch.set_default_tensor_type(torch.cuda.FloatTensor)
  print("using cuda:",torch.cuda.get_device_name(0))

hdf5_file = 'mount/My Drive/Colab Notebooks/celeba_dataset/celeba_aligned_small.h5py'
total_images = 10000
with h5py.File(hdf5_file, 'w') as hf:
    count = 0
    with zipfile.ZipFile('celeba/img_align_celeba.zip', 'r') as zf:
      for i in zf.namelist():
        if (i[-4:] == '.jpg'):
          ofile = zf.extract(i)
          img = imageio.imread(ofile)
          os.remove(ofile)
          hf.create_dataset('img_align_celeba/'+str(count)+'.jpg', data=img, compression="gzip", compression_opts=9)
          if (count == total_images):
            break
hdf5_file = 'mount/My Drive/Colab Notebooks/celeba_dataset/celeba_aligned_small.h5py'

import torch
import torch.nn as nn
from torch.utils.data import Dataset
import pandas
class CelebADataset(Dataset):
  def __init__(self,file):
    self.file_object = h5py.File(file,'r')
    self.dataset = self.file_object['img_align_celeba']

  def __len__(self):
    return len(self.dataset)

  def __getitem__(self,index):
    if(index>=len(self.dataset)):
      raise IndexError()
    img = numpy.array(self.dataset[str(index)+'.jpg'])
    return torch.cuda.FloatTensor(img)/255.0
  
  def plot_image(self,index):
    plt.imshow(numpy.array(self.dataset[str(index)+'.jpg']),interpolation='nearest')


def generate_random_image(size):
  random_data = torch.rand(size)
  return random_data
def generate_random_seed(size):
  random_data = torch.randn(size)
  return random_data
celeba_dataset = CelebADataset('mount/My Drive/Colab Notebooks/celeba_dataset/celeba_aligned_small.h5py')


class View(nn.Module):
  def __init__(self,shape):
    super().__init__()
    self.shape = shape,

  def forward(self,x):
    return x.view(*self.shape)

class Discriminator(nn.Module):
  def __init__(self):
    super().__init__()
    self.model = nn.Sequential(
        View(218*178*3),
        nn.Linear(3*218*178,100),
        nn.LeakyReLU(),
        nn.LayerNorm(100),
        nn.Linear(100,1),
        nn.Sigmoid()
    )
    self.loss_function = nn.BCELoss()
    self.optimiser = torch.optim.Adam(self.parameters(), lr=0.0001)
    self.counter = 0
    self.progress = []

  def forward(self, inputs):
    return self.model(inputs)

  def train(self, inputs, targets):
    outputs = self.forward(inputs)
    loss = self.loss_function(outputs, targets)
    self.optimiser.zero_grad()
    loss.backward()
    self.optimiser.step()

  def plot_progress(self):
    df = pandas.DataFrame(self.progress, columns=['loss'])
    df.plot(ylim=(0), figsize=(16,8), alpha=0.1, marker='.', grid=True, yticks=(0,0.25,1.0,0.5,5.0))

device=torch.device("cuda")
D = Discriminator()
D.to(device)
for image_data_tensor in celeba_dataset:
  D.train(image_data_tensor,torch.cuda.FloatTensor([1.0]))
  D.train(generate_random_image((218,178,3)),torch.cuda.FloatTensor([0.0]))

class Generator(nn.Module):
  def __init__(self):
    super().__init__()
    self.model = nn.Sequential(
        nn.Linear(100,3*10*10),
        nn.LeakyReLU(),
        nn.LayerNorm(3*10*10),
        nn.Linear(3*10*10,3*218*178),
        nn.Sigmoid(),
        View((218,178,3))
    )
    self.optimiser = torch.optim.Adam(self.parameters(), lr=0.0001)
    self.counter=0
    self.progress=[]
  def forward(self, inputs):
    return self.model(inputs)

  def train(self, D, inputs, targets):
    g_output = self.forward(inputs)
    d_output = D.forward(g_output)
    loss = D.loss_function(d_output, targets)
    self.optimiser.zero_grad()
    loss.backward()
    self.optimiser.step()
    
  def plot_progress(self):
    df = pandas.DataFrame(self.progress, columns=['loss'])
    df.plot(ylim=(0), figsize=(16,8), alpha=0.1, marker='.', grid=True, yticks=(0,0.25,1.0,0.5,5.0))

G = Generator()
G.to(device)
output = G.forward(generate_random_seed(100))
img = output.detach().cpu().numpy()
plt.imshow(img,interpolation='none',cmap='Blues')


D = Discriminator()
D.to(device)
G = Generator()
G.to(device)
epochs=6
for epoch in range(epochs):
  for image_data_tensor in celeba_dataset:
    D.train(image_data_tensor,torch.cuda.FloatTensor([1.0]))
    D.train(G.forward(generate_random_seed(100)).detach(),torch.cuda.FloatTensor([0.0])) 
    G.train(D,generate_random_seed(100),torch.cuda.FloatTensor([1.0]))

f,axarr = plt.subplots(2,3,figsize=(16,8))
for i in range(2):
  for j in range(3):
    output = G.forward(generate_random_seed(100))
    img = output.detach().cpu().numpy()
    axarr[i,j].imshow(img,interpolation='none',cmap='Blues')