生成對抗網路DCGAN

時間20210504

作者：知道許多的橘子

實作：生成對抗網路DCGAN_on_MNIST

如果感覺算力不夠用了，或者心疼自己電腦了！

可以用我實驗室的算力，試試呢！

害，誰叫我的算力都用不完呢！

支持所有框架！實際上框架都配置好了！

傻瓜式云計算！

Tesla v100 1卡，2卡，4卡，8卡

記憶體16-128G

cpu：8-24核

想要？加個微信：15615634293

歡迎打擾！

import torch
from torch import nn
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from torch import optim
import os

# 設定超引數
batch_size = 64
learning_rate = 0.0002
epochsize = 60
sample_dir = "images_3"

# 創建生成影像的目錄
if not os.path.exists(sample_dir):
    os.makedirs(sample_dir)


# 生成器結構
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        # 全連接層
        self.fc_layer = nn.Sequential(nn.Linear(100, 128 * 7 * 7))  # torch.Size([64, 6272])

        # 卷積層
        self.conv_layer = nn.Sequential(

            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2),  # torch.Size([64, 128, 14, 14])

            nn.Conv2d(128, 128, 3, stride=1, padding=1),
            nn.BatchNorm2d(128, 0.8),  
            
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),  # torch.Size([64, 128, 28, 28])

            nn.Conv2d(128, 64, 3, stride=1, padding=1),
            nn.BatchNorm2d(64, 0.8),
            nn.LeakyReLU(0.2, inplace=True),  # torch.Size([64, 64, 28, 28])

            nn.Conv2d(64, 1, 3, stride=1, padding=1),
            nn.Tanh()  # torch.Size([64, 1, 28, 28])
        )

    def forward(self, input):
        x = self.fc_layer(input)  # # torch.Size([64, 6272])
        x = x.view(input.shape[0], 128, 7, 7)  # torch.Size([64, 128, 7, 7])
        x = self.conv_layer(x)
        return x


# 鑒別器結構
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        # 卷積層
        self.conv_layer = nn.Sequential(

            # 第一層卷積層沒有使用bn層，其余三層均有使用
            nn.Conv2d(1, 16, 3, 2, 1),  # torch.Size([64, 16, 14, 14])
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout2d(0.25),

            nn.Conv2d(16, 32, 3, 2, 1),  # torch.Size([64, 32, 7, 7])
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout2d(0.25),
            nn.BatchNorm2d(32, momentum=0.8),

            nn.Conv2d(32, 64, 3, 2, 1),  # torch.Size([64, 64, 4, 4])
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout2d(0.25),
            nn.BatchNorm2d(64, momentum=0.8),

            nn.Conv2d(64, 128, 3, 2, 1),  # torch.Size([64, 128, 2, 2])
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout2d(0.25),
            nn.BatchNorm2d(128, momentum=0.8)
        )

        # 全連接層
        self.fc_layer = nn.Sequential(
            nn.Linear(512, 1),
            nn.Sigmoid()
        )

    def forward(self, input):
        x = self.conv_layer(input)  # torch.Size([64, 128, 2, 2])
        x = x.view(input.shape[0], -1)  # torch.Size([64, 512])
        x = self.fc_layer(x)  # torch.Size([64, 1])

        return x


# 訓練集下載
mnist_traindata = datasets.MNIST('/home/megstudio/dataset/dataset-2105/file-1258/mnist', train=True, transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5],
                         std=[0.5])
]), download=False)
mnist_train = DataLoader(mnist_traindata, batch_size=batch_size, shuffle=True)

# GPU加速
device = torch.device('cuda')

G = Generator().to(device)
D = Discriminator().to(device)

# 匯入之前的訓練模型
# G.load_state_dict(torch.load('G_plus.ckpt'))
# D.load_state_dict(torch.load('D_plus.ckpt'))

# 設定優化器與損失函式,二分類的時候使用BCELoss較好,BCEWithLogitsLoss是自帶一層Sigmoid
# criteon = nn.BCEWithLogitsLoss()
criteon = nn.BCELoss()
G_optimizer = optim.Adam(G.parameters(), lr=learning_rate)
D_optimizer = optim.Adam(D.parameters(), lr=learning_rate)

# 設定對比標簽
# realimage_label = torch.ones(batch_size, 1).to(device)  # value：1 torch.Size([128, 1])
# fakeimage_label = torch.zeros(batch_size, 1).to(device)  # value：0 torch.Size([128, 1])

# 開始訓練
print("start training")
for epoch in range(epochsize):

    D_loss_total = 0
    G_loss_total = 0
    total_num = 0

    # 這里的RealImageLabel是沒有用上的
    for batchidx, (realimage, _) in enumerate(mnist_train):

        realimage = realimage.to(device)
        realimage_label = torch.ones(realimage.size(0), 1).to(device)  # value：1 torch.Size([128, 1])
        fakeimage_label = torch.zeros(realimage.size(0), 1).to(device)  # value：0 torch.Size([128, 1])

        # 隨機生成噪音
        z = torch.randn(realimage.size(0), 100).to(device)

        #==================================================
        # 訓練鑒別器————總的損失為兩者相加
        d_realimage_loss = criteon(D(realimage), realimage_label) 
        d_fakeimage_loss = criteon(D(G(z)), fakeimage_label)  
        D_loss = d_realimage_loss + d_fakeimage_loss

        # 引數訓練三個步驟
        D_optimizer.zero_grad()
        D_loss.backward()
        D_optimizer.step()

        # 計算一次epoch的總損失
        D_loss_total += D_loss

        #==================================================
        # 訓練生成器————損失只有一個
        G_loss = criteon(D(G(z)), realimage_label)

        # 引數訓練三個步驟
        G_optimizer.zero_grad()
        G_loss.backward()
        G_optimizer.step()

        # 計算一次epoch的總損失
        G_loss_total += G_loss

        # 計算訓練影像個數
        total_num += realimage.size(0)

        # 列印相關的loss值
        if batchidx % 300 == 0:
            print("batchidx:{}/{}, D_loss:{}, G_loss:{}, total_num:{},".format(batchidx, len(mnist_train), D_loss, G_loss,
                                                                             total_num, ))

    # 列印一次訓練的loss值
    print('Epoch:{}/{}, D_loss:{}, G_loss:{}, total_num：{}'.format(epoch, epochsize, D_loss_total / len(mnist_train),
                                                                   G_loss_total / len(mnist_train), total_num))

    # 保存生成影像
    z = torch.randn(batch_size, 100).to(device)
    save_image(G(z).data[:64], os.path.join(sample_dir, 'fake_images-{}.png'.format(epoch + 14)), nrow=8,
               normalize=True)

    # 保存網路結構
    torch.save(G.state_dict(), 'G_plus.ckpt')
    torch.save(D.state_dict(), 'D_plus.ckpt')