gan训练

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms
from torchvision.utils import save_image
from torch.autograd import Variable
import os
 
if not os.path.exists('./img'):
    os.mkdir('./img')
 
 
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)
    out = out.view(-1, 1, 28, 28)
    return out
 
 
batch_size = 128
num_epoch = 100
z_dimension = 100
 
# Image processing
img_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])            #可能时mnist数据集数据更新过，以前使用的[0.5,0.5,0.5]会报错
])
# MNIST dataset
mnist = datasets.MNIST(
    root='./tensorflow/data_gan/', train=True, transform=img_transform, download=True)
# Data loader
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True)
 
 
# Discriminator
class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(784, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1), 
            nn.Sigmoid())
 
    def forward(self, x):
        x = self.dis(x)
        return x
 
 
# Generator
class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(True),
            nn.Linear(256, 256), 
            nn.ReLU(True), 
            nn.Linear(256, 784), 
            nn.Tanh())
 
    def forward(self, x):
        x = self.gen(x)
        return x
 
 
D = discriminator()
G = generator()
if torch.cuda.is_available():
    D = D.cuda()
    G = G.cuda()
# Binary cross entropy loss and optimizer
criterion = nn.BCELoss()
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)
 
# Start training
for epoch in range(num_epoch):
    for i, (img, _) in enumerate(dataloader):
        num_img = img.size(0)
        # =================train discriminator
        img = img.view(num_img, -1)
        real_img = Variable(img)
        real_label = Variable(torch.ones(num_img))
        fake_label = Variable(torch.zeros(num_img))
 
        # compute loss of real_img
        real_out = D(real_img)
        d_loss_real = criterion(real_out, real_label)
        real_scores = real_out  # closer to 1 means better
 
        # compute loss of fake_img
        z = Variable(torch.randn(num_img, z_dimension))
        fake_img = G(z)
        fake_out = D(fake_img)
        d_loss_fake = criterion(fake_out, fake_label)
        fake_scores = fake_out  # closer to 0 means better
 
        # bp and optimize
        d_loss = d_loss_real + d_loss_fake
        d_optimizer.zero_grad()
        d_loss.backward()
        d_optimizer.step()
 
        # ===============train generator
        # compute loss of fake_img
        z = Variable(torch.randn(num_img, z_dimension))
        fake_img = G(z)
        output = D(fake_img)
        g_loss = criterion(output, real_label)
 
        # bp and optimize
        g_optimizer.zero_grad()
        g_loss.backward()
        g_optimizer.step()
 
        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], d_loss: {:.6f}, g_loss: {:.6f} '
                  'D real: {:.6f}, D fake: {:.6f}'.format(
                      epoch, num_epoch, d_loss.item(), g_loss.item(),
                      real_scores.data.mean(), fake_scores.data.mean()))
    if epoch == 0:
        real_images = to_img(real_img.cpu().data)
        save_image(real_images, './img/real_images.png')
 
    fake_images = to_img(fake_img.cpu().data)
    save_image(fake_images, './img/fake_images-{}.png'.format(epoch + 1))
 
torch.save(G.state_dict(), './generator.pth')
torch.save(D.state_dict(), './discriminator.pth')

import torch
from torch import nn
from torch import autograd
from PIL import Image
from torchvision import transforms,utils

class discriminator(nn.Module):
    def __init__(self,inputsize,outputsize):
        super(discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(inputsize, outputsize),
            nn.LeakyReLU(0.2),
            nn.Linear(outputsize, outputsize),
            nn.LeakyReLU(0.2),
            nn.Linear(outputsize, 1),
            nn.Sigmoid()
        )
 
    def forward(self, x):
        x = self.dis(x)
        return x

class generator(nn.Module):
    def __init__(self,gsize,inputsize,outputsize):
        super(generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(gsize, outputsize),
            nn.ReLU(True),
            nn.Linear(outputsize, outputsize),
            nn.ReLU(True),
            nn.Linear(outputsize, inputsize),
            nn.Tanh()
        )
 
    def forward(self, x):
        x = self.gen(x)
        return x

width = 240
height = 360
x = width//8
y = height//12
z_dimension = 96

def to_img(img):
    out = 0.5 * (img + 1)
    out = out.clamp(0, 1)
    out = out.view(-1, 1, x, y)
    return out

image = Image.open('./tensorflow/jpg/text.png')
img = transforms.Compose([transforms.Resize((height,width)),
        transforms.Grayscale(),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])])(image)

print(img.shape)
img = torch.cat(img.chunk(12,1))#图像数据分割
img = torch.cat(img.chunk(8,2))#图像数据分割
img = img.view(12*8, -1)  # 将格子中图片展开
print(img.shape)

D = discriminator(x*y,8*12)
G = generator(z_dimension,x*y,8*12)

criterion = nn.BCELoss()
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0003)


epoch = 1000
for i in range(epoch*100):
    num_img = img.size(0)
    real_img = autograd.Variable(img)  # 将tensor变成Variable放入计算图中
    real_label = autograd.Variable(torch.ones(num_img))  # 定义真实label为1
    fake_label = autograd.Variable(torch.zeros(num_img))  # 定义假的label为0

    if i==0:
        utils.save_image(to_img(real_img.data), './tensorflow/gan-png/real_images.png')

    # compute loss of real_img
    real_out = D(real_img)  # 将真实的图片放入判别器中
    d_loss_real = criterion(real_out, real_label)  # 得到真实图片的loss  
    real_scores = real_out  # 真实图片放入判别器输出越接近1越好

    # compute loss of fake_img
    z = autograd.Variable(torch.randn(num_img, z_dimension))  # 随机生成一些噪声
    fake_img = G(z)  # 放入生成网络生成一张假的图片
    fake_out = D(fake_img)  # 判别器判断假的图片
    d_loss_fake = criterion(fake_out, fake_label)  # 得到假的图片的loss
    fake_scores = fake_out  # 假的图片放入判别器越接近0越好

    # bp and optimize
    d_loss = d_loss_real+d_loss_fake  # 将真假图片的loss加起来
    d_optimizer.zero_grad()  # 归0梯度
    d_loss.backward()  # 反向传播
    d_optimizer.step()  # 更新参数

    # compute loss of fake_img
    z = autograd.Variable(torch.randn(num_img, z_dimension))  # 得到随机噪声
    fake_img = G(z)  # 生成假的图片
    output = D(fake_img)  # 经过判别器得到结果
    g_loss = criterion(output, real_label)  # 得到假的图片与真实图片label的loss

    # bp and optimize
    g_optimizer.zero_grad()  # 归0梯度
    g_loss.backward()  # 反向传播
    g_optimizer.step()  # 更新生成网络的参数

    if i%(epoch) == 0:
        print("real_scores:%f,fake_scores:%f,d_loss:%f,g_loss:%f"%(real_scores.data.mean().item(),
        fake_scores.data.mean().item(),d_loss_fake.item(),g_loss.item()))
        print("-------------------------%d"%i)
        utils.save_image(to_img(fake_img.data), './tensorflow/gan-png/fake_images_{}.png'.format(i//epoch))