gan生成对抗网络

import argparse
import os
import random
import argparse
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML
from PIL import Image

manualSeed = 999
print("Random seed:", manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)

data_root = "./CAN_image_dataset/"  # dataset root
workers = 1  # using thread numbers
batch_size = 128  # batch_size
nc = 1  # number of channel from input images
num_epochs = 100  # number of training epochs
lr = 0.0001  # learning rate
beta1 = 0.5  # hyperparameter for adam optimizer
ngpu = 1

dataset = dset.ImageFolder(root=data_root,
                           transform=transforms.Compose([
                               transforms.Grayscale(1),
                               transforms.ToTensor(),
                           ]))

dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=workers)
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")

real_batch = next(iter(dataloader))


class PrintLayer(nn.Module):
    def __init__(self):
        super(PrintLayer, self).__init__()

    def forward(self, x):
        # Do your print / debug stuff here
        print(x)
        return x


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)


class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # Input : N x channel noise x 1 x 1
            nn.ConvTranspose2d(256, 512, (4, 3), stride=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            # second layer
            nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            # third layer
            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            # fourth layer,
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            # Final layer
            nn.ConvTranspose2d(64, 1, 4, stride=2, padding=1, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        return self.main(input)


netG = Generator(ngpu).to(device)

if (device.type == 'cuda') and (ngpu > 1):
    netG = nn.DataParallel(netG, list(range(ngpu)))

netG.apply(weights_init)

print(netG)


class Discriminator(nn.Module):
    def __init__(self, ngpu):
        super(Discriminator, self).__init__()
        self.ngpu = ngpu,
        self.main = nn.Sequential(
            nn.Conv2d(1, 1, (4, 3), stride=(2, 1), padding=1, bias=False),
            nn.BatchNorm2d(1),
            nn.ReLU(True),
            nn.Conv2d(1, 1, (4, 3), stride=(2, 1), padding=1, bias=False),
            nn.BatchNorm2d(1),
            nn.ReLU(True),
            nn.Conv2d(1, 1, (16, 48), stride=1, padding=0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)


netD = Discriminator(ngpu).to(device)

if (device.type == 'cuda') and (ngpu > 1):
    netD = nn.DataParallel(netD, list(range(ngpu)))

netD.apply(weights_init)

print(netD)

criterion = nn.BCELoss()
fixed_noise = torch.randn(64, 256, 1, 1, device=device)

real_label = 1
fake_label = 0

optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))

# train process0
from tqdm import tqdm
import time

# save losses to check training state
img_list = []
G_losses = []
D_losses = []
iters = 0

print("Starting Training Loop...")

# iterate epoch
for epoch in tqdm(range(num_epochs)):
    # iterate batch in an epoch
    for i, data in enumerate(dataloader, 0):
        ####################
        # (1) update discriminator network : maximize log(D(x)) + log(1 - D(G(z))))
        ####################
        # @@train real data@@
        netD.zero_grad()
        # Adapt to the size of the batches or the device to be used
        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0)
        label = torch.full((b_size,), real_label,
                           dtype=torch.float, device=device)
        # pass batch composed real data to D
        output = netD(real_cpu).view(-1)
        # get losses
        errD_real = criterion(output, label)
        # calculate degree of change while doing backpropagation
        errD_real.backward()
        D_x = output.mean().item()

        # @@ train fake data @@
        # Generate latent space vector used in Generator
        noise = torch.randn(b_size, 256, 1, 1, device=device)
        # Generate fake image using G
        fake = netG(noise)
        label.fill_(fake_label)
        # Discriminate truth of data using D
        output = netD(fake.detach()).view(-1)
        # Calculate losses of D
        errD_fake = criterion(output, label)
        # Calculate changes through backpropagation and accumulate changes which get before
        errD_fake.backward()
        D_G_z1 = output.mean().item()
        # add losses which get from both fake image and real image
        # At this time, errD is not used in backpropagation, but is used when reporting the learning state afterwards.
        errD = errD_fake + errD_real

        # Update D
        optimizerD.step()

        ###################
        # (2) Update G network : maximize log(D(G(z)))
        ###################
        netG.zero_grad()
        label.fill_(real_label)  # we use real label to get losses of Generator
        # Pass fake data to D again because we update D just now
        # At this time, G didn't update, but we get different value because D update
        output = netD(fake).view(-1)
        # get losses of G
        errG = criterion(output, label)
        # Calculate changes of G
        errG.backward()
        D_G_z2 = output.mean().item()
        # Update G
        optimizerG.step()

        # print training state
        if i % 50 == 0:
            print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                  % (epoch, num_epochs, i, len(dataloader),
                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))

        # save losses to draw graph later
        G_losses.append(errG.item())
        D_losses.append(errD.item())

        # save return value of G passed fixed noise
        if (iters % 500 == 0) or ((epoch == num_epochs - 1) and (i == len(dataloader) - 1)):
            with torch.no_grad():
                fake = netG(fixed_noise).detach().cpu()
            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

        iters += 1
#%%

plt.figure(figsize=(10,5))
plt.title('Generator and Discriminator Loss During Training')
plt.plot(G_losses, label='G')
plt.plot(D_losses, label='D')
plt.xlabel('iteration')
plt.ylabel('Loss')
plt.legend()
plt.show()

#%%

fig = plt.figure(figsize=(8,8))
plt.axis('off')
ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated = True)] for i in img_list]
ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000,blit=True)

HTML(ani.to_jshtml())

#%%

# bring real datas form dataloader
real_batch = next(iter(dataloader))

# print real images on display
plt.figure(figsize=(15,15))
plt.subplot(1,2,1)
plt.axis('off')
plt.title('Fake CAN Images')
plt.imshow(np.transpose(img_list[-1],(1,2,0)))
plt.show()

can_image

import numpy as np
import pandas as pd
from PIL import Image

"""
user setting 
1. write can_dataset_path which you download in reference paper
2. make folder to save can_image_dataset
3. run this code
"""

can_dataset_path = "dataset/Training set/normal_run_data.txt"


def one_hot_vector(c):
    # convert element of can_data to one hot vector to select can image position
    # can element is white(255) and else is black(0)
    #    can order: 1  2  3  4  5  6  7  8  9  0  a  b  c  d  e  f
    ohv = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    if c == "1":
        ohv[0] = 255
    elif c == "2":
        ohv[1] = 255
    elif c == "3":
        ohv[2] = 255
    elif c == "4":
        ohv[3] = 255
    elif c == "5":
        ohv[4] = 255
    elif c == "6":
        ohv[5] = 255
    elif c == "7":
        ohv[6] = 255
    elif c == "8":
        ohv[7] = 255
    elif c == "9":
        ohv[8] = 255
    elif c == "0":
        ohv[9] = 255
    elif c == "a":
        ohv[10] = 255
    elif c == "b":
        ohv[11] = 255
    elif c == "c":
        ohv[12] = 255
    elif c == "d":
        ohv[13] = 255
    elif c == "e":
        ohv[14] = 255
    elif c == "f":
        ohv[15] = 255
    return ohv


def convert_to_can_data_to_one_hot_vector(can):
    # convert can_data to image

    c1 = one_hot_vector(can[0])
    c2 = one_hot_vector(can[1])
    c3 = one_hot_vector(can[2])

    image_can = np.concatenate((c1,c2,c3),axis=0)
    return image_can


def make_can_image(data):
    row_size = 64
    column_size = 48

    # set size of image to 16 X ( 3 * 64 )
    # colum is Hexadecimal presentation from CAN ID value
    # row is  3 (CAN ID length) * 64 (optimal row size in reference paper)

    can_img_num = 0
    for w in range(0, len(data), row_size):
        can_image = np.zeros((1, column_size))
        lst = data.iloc[w : w + row_size]

        for i in range(0, len(lst)):
            add_can = convert_to_can_data_to_one_hot_vector(lst.iloc[i])
            can_image = np.vstack([can_image, add_can])
        can_image = np.delete(can_image, (0), axis=0)

        # convert (can_img => int type can_img => RGB can_img) to save to png file form
        int_can_image = can_image.astype(int)
        image = Image.fromarray(np.uint8(int_can_image))

        # save img file
        image.save("CAN_image_dataset/Training_set/can_img_{}.png".format(can_img_num))
        can_img_num += 1

    print("making can image complete | the number of images : ", can_img_num)


# test this code to check working well


if __name__ == "__main__":
    """load normal_run_data.txt"""
    can_feature = ["Timestamp", "blank", "ID", "zero", "DLC", "Data"]
    data_for_test = pd.read_csv(
        can_dataset_path,
        sep="    ",
        engine="python",
        encoding="cp949",
        header=None,
        names=can_feature,
    )
    df = data_for_test["ID"]
    df = df.drop(df.index[988871])  # drop None index

    for i in range(len(df)):  # extract needed part
        df.at[i] = df.at[i][5:8]

    make_can_image(df)

    # use this code later to load can image
    # img = Image.open('filename.png').convert('L')
    # img_array = np.array(img)
    # print(img_array)