Pytorch实现图片异常检测功能

import os
import urllib.request
import zipfile
import tarfile

import matplotlib.pyplot as plt
%matplotlib inline
from PIL import Image
import numpy as np

#不存在“data”文件夹时创建
data_dir = "./data/"
if not os.path.exists(data_dir):
    os.mkdir(data_dir)

import sklearn
# 下载并读取MNIST的手写数字图像。
from sklearn.datasets import fetch_openml

mnist = fetch_openml('mnist_784', version=1, data_home="./data/")  #data_home指定保存地址。

# 数据的提取
X = mnist.data
y = mnist.target

# 将MNIST的第一个数据可视化
plt.imshow(np.array(X.iloc[0]).reshape(28, 28), cmap='gray')
print("这个图像数据的标签是{}".format(y[0]))

#在文件夹“data”下创建文件夹“img_78”
data_dir_path = "./data/img_78/"
if not os.path.exists(data_dir_path):
    os.mkdir(data_dir_path)

#从MNIST将数字7、8的图像作为图像保存到“img_78”文件夹中
count7=0
count8=0
max_num=200  # 每制作200张图片

for i in range(len(X)):
    
    # 图像7的制作
    if (y[i] == "7") and (count7<max_num):
        file_path="./data/img_78/img_7_"+str(count7)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28×28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f = pil_img_f.resize((64, 64), Image.BICUBIC)  # 扩大到64×64
        pil_img_f.save(file_path)  # 保存
        count7+=1 
    
    #图像8的制作
    if (y[i] == "8") and (count8<max_num):
        file_path="./data/img_78/img_8_"+str(count8)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28*28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f = pil_img_f.resize((64, 64), Image.BICUBIC)  # 扩大到64×64
        pil_img_f.save(file_path)  # 保存
        count8+=1
        
    # 制作200张7和8之后，break
    if (count7>=max_num) and (count8>=max_num):
        break


# 在文件夹“data”下面创建文件夹“test”
data_dir_path = "./data/test/"
if not os.path.exists(data_dir_path):
    os.mkdir(data_dir_path)

# 在上述制作7,8图像时使用的index的最终值
i_start = i+1
print(i_start)

# 从MNIST将数字7、8的图像作为图像保存到“img_78”文件夹中
count2=0
count7=0
count8=0
max_num=5  #每制作五张图片

for i in range(i_start,len(X)):  # 从i_start开始
    
    #图像2的制作
    if (y[i] == "2") and (count2<max_num):
        file_path="./data/test/img_2_"+str(count2)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28×28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f = pil_img_f.resize((64, 64), Image.BICUBIC)  # 扩大到64×64
        pil_img_f.save(file_path)  # 保存
        count2+=1
    
    # 图像7的制作
    if (y[i] == "7") and (count7<max_num):
        file_path="./data/test/img_7_"+str(count7)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  #将图像变形为28×28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f = pil_img_f.resize((64, 64), Image.BICUBIC)  # 6扩大到64×64
        pil_img_f.save(file_path)  # 保存
        count7+=1 
    
    # 图像8的制作
    if (y[i] == "8") and (count8<max_num):
        file_path="./data/test/img_8_"+str(count8)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28*28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f = pil_img_f.resize((64, 64), Image.BICUBIC)  # 扩大到64×64
        pil_img_f.save(file_path)  # 保存
        count8+=1 


# 在文件夹“data”下创建文件夹“img_78_28size”
data_dir_path = "./data/img_78_28size/"
if not os.path.exists(data_dir_path):
    os.mkdir(data_dir_path)

# 从MNIST将数字7、8的图像作为图像保存到“img_78_28size”文件夹中
count7=0
count8=0
max_num=200  # 每制作200张图片

for i in range(len(X)):
    
    # 图像7的制作
    if (y[i] == "7") and (count7<max_num):
        file_path="./data/img_78_28size/img_7_"+str(count7)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28×28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f.save(file_path)  # 保存
        count7+=1 
    
    # 图像8的制作
    if (y[i] == "8") and (count8<max_num):
        file_path="./data/img_78_28size/img_8_"+str(count8)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28*28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 画像变成PIL
        pil_img_f.save(file_path)  # 保存
        count8+=1
    
    if (count7>=max_num) and (count8>=max_num):
        break



# 在文件夹“data”下面创建文件夹“test”
data_dir_path = "./data/test_28size/"
if not os.path.exists(data_dir_path):
    os.mkdir(data_dir_path)

# 在上述制作7,8图像时使用的index的最终值
i_start = i+1
print(i_start)

# 从MNIST将数字7、8的图像作为图像保存到“img_78”文件夹中
count2=0
count7=0
count8=0
max_num=5  # 每制作五张图片

for i in range(i_start,len(X)):  #从i_start开始
    
    # 图像2的制作
    if (y[i] == "2") and (count2<max_num):
        file_path="./data/test_28size/img_2_"+str(count2)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28×28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f.save(file_path)  # 保存
        count2+=1 
    
    # 画像7的制作
    if (y[i] == "7") and (count7<max_num):
        file_path="./data/test_28size/img_7_"+str(count7)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28×28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f.save(file_path)  # 保存
        count7+=1 
    
    # 图像8的制作
    if (y[i] == "8") and (count8<max_num):
        file_path="./data/test_28size/img_8_"+str(count8)+".jpg"
        im_f=(np.array(X.iloc[i]).reshape(28, 28))  # 将图像变形为28*28
        pil_img_f = Image.fromarray(im_f.astype(np.uint8))  # 把图像变成PIL
        pil_img_f.save(file_path)  # 保存
        count8+=1 

# 导入软件包
import random
import math
import time
import pandas as pd
import numpy as np
from PIL import Image

import torch
import torch.utils.data as data
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torchvision import transforms

# Setup seeds
torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)

class Generator(nn.Module):

    def __init__(self, z_dim=20, image_size=64):
        super(Generator, self).__init__()

        self.layer1 = nn.Sequential(
            nn.ConvTranspose2d(z_dim, image_size * 8,
                               kernel_size=4, stride=1),
            nn.BatchNorm2d(image_size * 8),
            nn.ReLU(inplace=True))

        self.layer2 = nn.Sequential(
            nn.ConvTranspose2d(image_size * 8, image_size * 4,
                               kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(image_size * 4),
            nn.ReLU(inplace=True))

        self.layer3 = nn.Sequential(
            nn.ConvTranspose2d(image_size * 4, image_size * 2,
                               kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(image_size * 2),
            nn.ReLU(inplace=True))

        self.layer4 = nn.Sequential(
            nn.ConvTranspose2d(image_size * 2, image_size,
                               kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(image_size),
            nn.ReLU(inplace=True))

        self.last = nn.Sequential(
            nn.ConvTranspose2d(image_size, 1, kernel_size=4,
                               stride=2, padding=1),
            nn.Tanh())
        #注意 ：由于是黑白图像，因此输出通道数量为1

    def forward(self, z):
        out = self.layer1(z)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.last(out)

        return out

# 动作确认
import matplotlib.pyplot as plt
%matplotlib inline

G = Generator(z_dim=20, image_size=64)

# 输入的随机数
input_z = torch.randn(1, 20)

# 将张量尺寸变形为(1,20,1,1)
input_z = input_z.view(input_z.size(0), input_z.size(1), 1, 1)

# 输出假图像
fake_images = G(input_z)

img_transformed = fake_images[0][0].detach().numpy()
plt.imshow(img_transformed, 'gray')
plt.show()

class Discriminator(nn.Module):

    def __init__(self, z_dim=20, image_size=64):
        super(Discriminator, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Conv2d(1, image_size, kernel_size=4,
                      stride=2, padding=1),
            nn.LeakyReLU(0.1, inplace=True))
        #注意 ：由于是黑白图像，因此输出通道数量为1

        self.layer2 = nn.Sequential(
            nn.Conv2d(image_size, image_size*2, kernel_size=4,
                      stride=2, padding=1),
            nn.LeakyReLU(0.1, inplace=True))

        self.layer3 = nn.Sequential(
            nn.Conv2d(image_size*2, image_size*4, kernel_size=4,
                      stride=2, padding=1),
            nn.LeakyReLU(0.1, inplace=True))

        self.layer4 = nn.Sequential(
            nn.Conv2d(image_size*4, image_size*8, kernel_size=4,
                      stride=2, padding=1),
            nn.LeakyReLU(0.1, inplace=True))

        self.last = nn.Conv2d(image_size*8, 1, kernel_size=4, stride=1)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)

        feature = out  #最后将通道集中到一个特征量中
        feature = feature.view(feature.size()[0], -1)  #转换为二维

        out = self.last(out)

        return out, feature

# 确认程序执
D = Discriminator(z_dim=20, image_size=64)

#生成伪造图像
input_z = torch.randn(1, 20)
input_z = input_z.view(input_z.size(0), input_z.size(1), 1, 1)
fake_images = G(input_z)

#将伪造的图像输入判别器D中
d_out = D(fake_images)

# 将输出值d_out乘以Sigmoid函数，将其转换成0～1的值
print(nn.Sigmoid()(d_out[0]))

# feature
print(d_out[1].shape)

def make_datapath_list():
    """制作用于学习、验证的图像数据和标注数据的文件路径表。 """

    train_img_list = list()  # 保存图像文件路径

    for img_idx in range(200):
        img_path = "./data/img_78/img_7_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

        img_path = "./data/img_78/img_8_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

    return train_img_list

class ImageTransform():
    """图像的预处理类"""

    def __init__(self, mean, std):
        self.data_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])

    def __call__(self, img):
        return self.data_transform(img)
class GAN_Img_Dataset(data.Dataset):
    """图像的Dataset类。继承PyTorch的Dataset类"""

    def __init__(self, file_list, transform):
        self.file_list = file_list
        self.transform = transform

    def __len__(self):
        '''返回图像的张数'''
        return len(self.file_list)

    def __getitem__(self, index):
        '''获取经过预处理后的图像的张量格式的数据'''

        img_path = self.file_list[index]
        img = Image.open(img_path)  # [高][宽]黑白

        #图像的预处理
        img_transformed = self.transform(img)

        return img_transformed

#创建DataLoader并确认操作

#创建文件列表
train_img_list=make_datapath_list()

# 创建Dataset
mean = (0.5,)
std = (0.5,)
train_dataset = GAN_Img_Dataset(
    file_list=train_img_list, transform=ImageTransform(mean, std))

# 创建DataLoader
batch_size = 64

train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)

#确认执行结果
batch_iterator = iter(train_dataloader)  # 转换成迭代器
imges = next(batch_iterator)   #取出位于第一位的元素
print(imges.size())  # torch.Size([64, 1, 64, 64])

#网络的初始化
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        #Conv2d和ConvTranspose2d的初始化
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
    elif classname.find('BatchNorm') != -1:
        # BatchNorm2d的初始化
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)


# 初始化的实施
G.apply(weights_init)
D.apply(weights_init)

print("网络已经成功地完成了初始化")

# 创建一个函数来学习模型


def train_model(G, D, dataloader, num_epochs):

    #确认是否可以使用GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("使用设备：", device)

    # 优化方法的设定
    g_lr, d_lr = 0.0001, 0.0004
    beta1, beta2 = 0.0, 0.9
    g_optimizer = torch.optim.Adam(G.parameters(), g_lr, [beta1, beta2])
    d_optimizer = torch.optim.Adam(D.parameters(), d_lr, [beta1, beta2])

    # 定义误差函数
    criterion = nn.BCEWithLogitsLoss(reduction='mean')

    # 使用硬编码的参数
    z_dim = 20
    mini_batch_size = 64

    #将网络变成GPU
    G.to(device)
    D.to(device)

    G.train()  # 将模型转换为训练模式
    D.train()  # 将模型转换为训练模式

    #如果网络相对固定，则开启加速
    torch.backends.cudnn.benchmark = True

    # 图像的张数
    num_train_imgs = len(dataloader.dataset)
    batch_size = dataloader.batch_size

    # 设置了迭代计数器
    iteration = 1
    logs = []

    #epoch循环
    for epoch in range(num_epochs):

        # 保存开始时间
        t_epoch_start = time.time()
        epoch_g_loss = 0.0  # epoch损失总和
        epoch_d_loss = 0.0  # epoch损失总和

        print('-------------')
        print('Epoch {}/{}'.format(epoch, num_epochs))
        print('-------------')
        print('（train）')

        # 以minibatch为单位从数据加载器中读取数据的循环
        for imges in dataloader:

            # --------------------
            # 1. 判别器D的学习
            # --------------------
            # 如果小批次的尺寸设置为1，会导致批次归一化处理产生错误，因此需要避免
            if imges.size()[0] == 1:
                continue

             #如果能使用GPU，则将数据送入GPU中
            imges = imges.to(device)

            #创建正确答案标签和伪造数据标签
           #在epoch最后的迭代中，小批次的数量会减少
            mini_batch_size = imges.size()[0]
            label_real = torch.full((mini_batch_size,), 1).to(device)
            label_fake = torch.full((mini_batch_size,), 0).to(device)

           #对真正的图像进行判定
            d_out_real, _ = D(imges)

             #生成伪造图像并进行判定
            input_z = torch.randn(mini_batch_size, z_dim).to(device)
            input_z = input_z.view(input_z.size(0), input_z.size(1), 1, 1)
            fake_images = G(input_z)
            d_out_fake, _ = D(fake_images)

            #计算误差
            d_loss_real = criterion(d_out_real.view(-1), label_real.float())
            d_loss_fake = criterion(d_out_fake.view(-1), label_fake.float())
            d_loss = d_loss_real + d_loss_fake

            #反向传播处理
            g_optimizer.zero_grad()
            d_optimizer.zero_grad()

            d_loss.backward()
            d_optimizer.step()

            # --------------------
            # 2.生成器G的学习
            # --------------------
           #生成伪造图像并进行判定
            input_z = torch.randn(mini_batch_size, z_dim).to(device)
            input_z = input_z.view(input_z.size(0), input_z.size(1), 1, 1)
            fake_images = G(input_z)
            d_out_fake, _ = D(fake_images)

            #计算误差
            g_loss = criterion(d_out_fake.view(-1), label_real.float())

           #反向传播处理
            g_optimizer.zero_grad()
            d_optimizer.zero_grad()
            g_loss.backward()
            g_optimizer.step()

            # --------------------
            # 3. 记录结果
            # --------------------
            epoch_d_loss += d_loss.item()
            epoch_g_loss += g_loss.item()
            iteration += 1

        #epoch的每个phase的loss和准确率
        t_epoch_finish = time.time()
        print('-------------')
        print('epoch {} || Epoch_D_Loss:{:.4f} ||Epoch_G_Loss:{:.4f}'.format(
            epoch, epoch_d_loss/batch_size, epoch_g_loss/batch_size))
        print('timer:  {:.4f} sec.'.format(t_epoch_finish - t_epoch_start))
        t_epoch_start = time.time()

    
    print("总迭代次数:", iteration)

    return G, D

# 进行训练和验证
num_epochs = 300
G_update, D_update = train_model(
    G, D, dataloader=train_dataloader, num_epochs=num_epochs)

# 将生成图像和训练数据可视化

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 输入的随机数生成
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_update(fixed_z.to(device))

# 训练数据
batch_iterator = iter(train_dataloader) #转换成迭代器
imges = next(batch_iterator) #提取第1个要素


#图像的可视化处理
fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
    #上层显示测试图像
    plt.subplot(2, 5, i+1)
    plt.imshow(imges[i][0].cpu().detach().numpy(), 'gray')

    #下层显示生成图像
    plt.subplot(2, 5, 5+i+1)
    plt.imshow(fake_images[i][0].cpu().detach().numpy(), 'gray')
    
def Anomaly_score(x, fake_img, D, Lambda=0.1):

    #求测试图像x和生成图像fake_img的像素级差的绝对值，并对每个迷你批求和
    residual_loss = torch.abs(x-fake_img)
    residual_loss = residual_loss.view(residual_loss.size()[0], -1)
    residual_loss = torch.sum(residual_loss, dim=1)

    # 将测试图像x和生成图像fake_img输入到识别器D，取出特征量
    _, x_feature = D(x)
    _, G_feature = D(fake_img)

    # 求测试图像x和生成图像fake_img的特征量之差的绝对值，对每个迷你批次求和
    discrimination_loss = torch.abs(x_feature-G_feature)
    discrimination_loss = discrimination_loss.view(
        discrimination_loss.size()[0], -1)
    discrimination_loss = torch.sum(discrimination_loss, dim=1)

    # 将两种损失对每个迷你批进行加法运算
    loss_each = (1-Lambda)*residual_loss + Lambda*discrimination_loss

    #求迷你批的全部损失
    total_loss = torch.sum(loss_each)

    return total_loss, loss_each, residual_loss

# 创建测试用的DataLoader


def make_test_datapath_list():
    """制作用于学习、验证的图像数据和标注数据的文件路径表。 """

    train_img_list = list()  # 保存图像文件路径

    for img_idx in range(5):
        img_path = "./data/test/img_7_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

        img_path = "./data/test/img_8_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

        img_path = "./data/test/img_2_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

    return train_img_list


# 制作文件列表
test_img_list = make_test_datapath_list()

# 制作Dataset
mean = (0.5,)
std = (0.5,)
test_dataset = GAN_Img_Dataset(
    file_list=test_img_list, transform=ImageTransform(mean, std))

# 制作DataLoader
batch_size = 5

test_dataloader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False)

# 测试数据的确认
batch_iterator = iter(test_dataloader)  # 转换成迭代器
imges = next(batch_iterator)  

# 取出第一个迷你批次

fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
    plt.subplot(2, 5, i+1)
    plt.imshow(imges[i][0].cpu().detach().numpy(), 'gray')

# 想检测异常的图像
x = imges[0:5]
x = x.to(device)

# 用于生成想要异常检测的图像的初始随机数
z = torch.randn(5, 20).to(device)
z = z.view(z.size(0), z.size(1), 1, 1)

# 変将requires_grad设为True，使得变量z可以求导数
z.requires_grad = True

#求z的优化函数，以便能够更新变量z
z_optimizer = torch.optim.Adam([z], lr=1e-3)


#求z
for epoch in range(5000+1):
    fake_img = G_update(z)
    loss, _, _ = Anomaly_score(x, fake_img, D_update, Lambda=0.1)

    z_optimizer.zero_grad()
    loss.backward()
    z_optimizer.step()

    if epoch % 1000 == 0:
        print('epoch {} || loss_total:{:.0f} '.format(epoch, loss.item()))

# 生成图像
fake_img = G_update(z)

# 要求损失
loss, loss_each, residual_loss_each = Anomaly_score(
    x, fake_img, D_update, Lambda=0.1)

#损失的计算总损失
loss_each = loss_each.cpu().detach().numpy()
print("total loss：", np.round(loss_each, 0))

# 图像可视化
fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
    # 把测试数据放在上层
    plt.subplot(2, 5, i+1)
    plt.imshow(imges[i][0].cpu().detach().numpy(), 'gray')

    # 在下层显示生成数据
    plt.subplot(2, 5, 5+i+1)
    plt.imshow(fake_img[i][0].cpu().detach().numpy(), 'gray')

# 导入软件包
import random
import math
import time
import pandas as pd
import numpy as np
from PIL import Image

import torch
import torch.utils.data as data
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torchvision import transforms

# Setup seeds
torch.manual_seed(1234)
torch.cuda.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)

class Generator(nn.Module):

    def __init__(self, z_dim=20):
        super(Generator, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Linear(z_dim, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(inplace=True))

        self.layer2 = nn.Sequential(
            nn.Linear(1024, 7*7*128),
            nn.BatchNorm1d(7*7*128),
            nn.ReLU(inplace=True))

        self.layer3 = nn.Sequential(
            nn.ConvTranspose2d(in_channels=128, out_channels=64,
                               kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True))

        self.last = nn.Sequential(
            nn.ConvTranspose2d(in_channels=64, out_channels=1,
                               kernel_size=4, stride=2, padding=1),
            nn.Tanh())
        #注意 ：由于是黑白图像，因此输出通道数量为 1

    def forward(self, z):
        out = self.layer1(z)
        out = self.layer2(out)

        #为了能置入卷积层中，需要对张量进行变形
        out = out.view(z.shape[0], 128, 7, 7)
        out = self.layer3(out)
        out = self.last(out)

        return out

#确认执行结果
import matplotlib.pyplot as plt
%matplotlib inline

G = Generator(z_dim=20)
G.train()

#输入的随机数
#由于要进行批次归一化处理，因此将小批次数设置为 2 以上
input_z = torch.randn(2, 20)

#输出伪造图像
fake_images = G(input_z)  # torch.Size([2, 1, 28, 28])
img_transformed = fake_images[0][0].detach().numpy()
plt.imshow(img_transformed, 'gray')
plt.show()

class Discriminator(nn.Module):

    def __init__(self, z_dim=20):
        super(Discriminator, self).__init__()

        #图像这边的输入处理
        self.x_layer1 = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=4,
                      stride=2, padding=1),
            nn.LeakyReLU(0.1, inplace=True))
      #注意 ：由于是黑白图像，因此输入通道数量为 1

        self.x_layer2 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=4,
                      stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.1, inplace=True))

      #随机数这边的输入处理
        self.z_layer1 = nn.Linear(z_dim, 512)

        #最终的判定
        self.last1 = nn.Sequential(
            nn.Linear(3648, 1024),
            nn.LeakyReLU(0.1, inplace=True))

        self.last2 = nn.Linear(1024, 1)

    def forward(self, x, z):

        #图像这边的输入处理
        x_out = self.x_layer1(x)
        x_out = self.x_layer2(x_out)

       #随机数这边的输入处理
        z = z.view(z.shape[0], -1)
        z_out = self.z_layer1(z)

        #将x_out与z_out连接在一起，交给全连接层进行判定
        x_out = x_out.view(-1, 64 * 7 * 7)
        out = torch.cat([x_out, z_out], dim=1)
        out = self.last1(out)

        feature = out  #最后将通道集中到一个特征量中
        feature = feature.view(feature.size()[0], -1)   #转换为二维

        out = self.last2(out)

        return out, feature


#确认执行结果
D = Discriminator(z_dim=20)

#生成伪造图像
input_z = torch.randn(2, 20)
fake_images = G(input_z)

#将伪造图像输入判定器D中
d_out, _ = D(fake_images, input_z)

#将输出结果d_out乘以Sigmoid，以将其转换为0～1的值
print(nn.Sigmoid()(d_out))


class Encoder(nn.Module):

    def __init__(self, z_dim=20):
        super(Encoder, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3,
                      stride=1),
            nn.LeakyReLU(0.1, inplace=True))
        #把图像转换成z

        self.layer2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3,
                      stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.1, inplace=True))

        self.layer3 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3,
                      stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True))

        #到这里为止，图像的尺寸为7像素×7像素
        self.last = nn.Linear(128 * 7 * 7, z_dim)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)

       #为了能放入FC中，对张量进行变形
        out = out.view(-1, 128 * 7 * 7)
        out = self.last(out)

        return out


#确认执行结果
E = Encoder(z_dim=20)

#输入的图像数据
x = fake_images  #fake_images是由上面的生成器G生成的

#将图像编码为z
z = E(x)

print(z.shape)
print(z)

def make_datapath_list():
    """制作用于学习、验证的图像数据和标注数据的文件路径表。 """

    train_img_list = list()  # 保存图像文件路径

    for img_idx in range(200):
        img_path = "./data/img_78_28size/img_7_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

        img_path = "./data/img_78_28size/img_8_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

    return train_img_list

class ImageTransform():
    """图像的预处理类"""

    def __init__(self, mean, std):
        self.data_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])

    def __call__(self, img):
        return self.data_transform(img)

class GAN_Img_Dataset(data.Dataset):
    """图像的Dataset类。继承PyTorch的Dataset类"""

    def __init__(self, file_list, transform):
        self.file_list = file_list
        self.transform = transform

    def __len__(self):
        '''返回图像的张数'''
        return len(self.file_list)

    def __getitem__(self, index):
        '''获取预处理图像的Tensor格式数据'''

        img_path = self.file_list[index]
        img = Image.open(img_path)  # [高][宽]黑白

        # 图像的预处理
        img_transformed = self.transform(img)

        return img_transformed

# 创建DataLoader并确认操作

#制作文件列表
train_img_list=make_datapath_list()

# Datasetを作成
mean = (0.5,)
std = (0.5,)
train_dataset = GAN_Img_Dataset(
    file_list=train_img_list, transform=ImageTransform(mean, std))

# 制作DataLoader
batch_size = 64

train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)

# 动作的确认
batch_iterator = iter(train_dataloader)  # 转换成迭代器
imges = next(batch_iterator)  # 找出第一个要素
print(imges.size())  # torch.Size([64, 1, 64, 64])

#创建用于训练模型的函数


def train_model(G, D, E, dataloader, num_epochs):

    #确认是否可以使用GPU加速
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("使用设备：", device)

    #设置最优化算法
    lr_ge = 0.0001
    lr_d = 0.0001/4
    beta1, beta2 = 0.5, 0.999
    g_optimizer = torch.optim.Adam(G.parameters(), lr_ge, [beta1, beta2])
    e_optimizer = torch.optim.Adam(E.parameters(), lr_ge, [beta1, beta2])
    d_optimizer = torch.optim.Adam(D.parameters(), lr_d, [beta1, beta2])

    #定义误差函数
    #BCEWithLogitsLoss是先将输入数据乘以Logistic，
    # 再计算二进制交叉熵
    criterion = nn.BCEWithLogitsLoss(reduction='mean')

    #对参数进行硬编码
    z_dim = 20
    mini_batch_size = 64

    #将网络载入GPU中
    G.to(device)
    E.to(device)
    D.to(device)

    G.train()  #将模型设置为训练模式
    E.train()  #将模型设置为训练模式
    D.train()  #将模型设置为训练模式

    #如果网络相对固定，则开启加速
    torch.backends.cudnn.benchmark = True

    #图像的张数
    num_train_imgs = len(dataloader.dataset)
    batch_size = dataloader.batch_size

    #设置迭代计数器
    iteration = 1
    logs = []

    # epoch循环
    for epoch in range(num_epochs):

        #保存开始时间
        t_epoch_start = time.time()
        epoch_g_loss = 0.0  #epoch的损失总和
        epoch_e_loss = 0.0  #epoch的损失总和
        epoch_d_loss = 0.0  #epoch的损失总和

        print('-------------')
        print('Epoch {}/{}'.format(epoch, num_epochs))
        print('-------------')
        print('（train）')

        #以minibatch为单位从数据加载器中读取数据的循环
        for imges in dataloader:

            #如果小批次的尺寸设置为1，会导致批次归一化处理产生错误，因此需要避免
            if imges.size()[0] == 1:
                continue

            #创建用于表示小批次尺寸为1和0的标签
            #创建正确答案标签和伪造数据标签
            #在epoch最后的迭代中，小批次的数量会减少
            mini_batch_size = imges.size()[0]
            label_real = torch.full((mini_batch_size,), 1).to(device)
            label_fake = torch.full((mini_batch_size,), 0).to(device)

            #如果能使用GPU，则将数据送入GPU中
            imges = imges.to(device)

            # --------------------
            # 1. 判别器D的学习
            # --------------------
            # 对真实的图像进行判定　
            z_out_real = E(imges)
            d_out_real, _ = D(imges, z_out_real)

            # 生成伪造图像并进行判定
            input_z = torch.randn(mini_batch_size, z_dim).to(device)
            fake_images = G(input_z)
            d_out_fake, _ = D(fake_images, input_z)

            #计算误差
            d_loss_real = criterion(d_out_real.view(-1), label_real.float())
            d_loss_fake = criterion(d_out_fake.view(-1), label_fake.float())
            d_loss = d_loss_real + d_loss_fake

            #反向传播
            d_optimizer.zero_grad()
            d_loss.backward()
            d_optimizer.step()

            # --------------------
            # 2. 生成器G的学习
            # --------------------
            #生成伪造图像并进行判定
            input_z = torch.randn(mini_batch_size, z_dim).to(device)
            fake_images = G(input_z)
            d_out_fake, _ = D(fake_images, input_z)

            #计算误差
            g_loss = criterion(d_out_fake.view(-1), label_real.float())

            #反向传播
            g_optimizer.zero_grad()
            g_loss.backward()
            g_optimizer.step()

            # --------------------
            # 3. 编码器E的学习
            # --------------------
            #对真实图像的z进行推定
            z_out_real = E(imges)
            d_out_real, _ = D(imges, z_out_real)

            #计算误差
            e_loss = criterion(d_out_real.view(-1), label_fake.float())

            #反向传播
            e_optimizer.zero_grad()
            e_loss.backward()
            e_optimizer.step()

            # --------------------
            #4.记录
            # --------------------
            epoch_d_loss += d_loss.item()
            epoch_g_loss += g_loss.item()
            epoch_e_loss += e_loss.item()
            iteration += 1

        #epoch的每个phase的loss和准确率
        t_epoch_finish = time.time()
        print('-------------')
        print('epoch {} || Epoch_D_Loss:{:.4f} ||Epoch_G_Loss:{:.4f} ||Epoch_E_Loss:{:.4f}'.format(
            epoch, epoch_d_loss/batch_size, epoch_g_loss/batch_size, epoch_e_loss/batch_size))
        print('timer:  {:.4f} sec.'.format(t_epoch_finish - t_epoch_start))
        t_epoch_start = time.time()

    print("总迭代次数:", iteration)

    return G, D, E

#网络的初始化
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
       #Conv2d和ConvTranspose2d的初始化
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
    elif classname.find('BatchNorm') != -1:
       # BatchNorm2d的初始化
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
    elif classname.find('Linear') != -1:
        #全连接层Linear的初始化
        m.bias.data.fill_(0)


#开始初始化
G.apply(weights_init)
E.apply(weights_init)
D.apply(weights_init)

print("网络已经成功地完成了初始化")

# 进行训练和验证
num_epochs = 1500
G_update, D_update, E_update = train_model(
    G, D, E, dataloader=train_dataloader, num_epochs=num_epochs)

#对生成图像与训练数据的可视化处理
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#生成输入的随机数
batch_size = 8
z_dim = 20
fixed_z = torch.randn(batch_size, z_dim)
fake_images = G_update(fixed_z.to(device))

#训练数据
batch_iterator = iter(train_dataloader)  #转换成迭代器
imges = next(batch_iterator)  #取出最开头的元素


#输出
fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
    #在上层中显示训练数据
    plt.subplot(2, 5, i+1)
    plt.imshow(imges[i][0].cpu().detach().numpy(), 'gray')

    #在下层中显示生成数据
    plt.subplot(2, 5, 5+i+1)
    plt.imshow(fake_images[i][0].cpu().detach().numpy(), 'gray')

# ·制作测试用的Dataloader


def make_test_datapath_list():
    """制作用于学习、验证的图像数据和标注数据的文件路径表。 """

    train_img_list = list()  # ·保存图像文件路径

    for img_idx in range(5):
        img_path = "./data/test_28size/img_7_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

        img_path = "./data/test_28size/img_8_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

        img_path = "./data/test_28size/img_2_" + str(img_idx)+'.jpg'
        train_img_list.append(img_path)

    return train_img_list


#制作文件列表
test_img_list = make_test_datapath_list()

# 创建Dataset
mean = (0.5,)
std = (0.5,)
test_dataset = GAN_Img_Dataset(
    file_list=test_img_list, transform=ImageTransform(mean, std))

# 制作DataLoader
batch_size = 5

test_dataloader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False)

#训练数据
batch_iterator = iter(test_dataloader)  #转换成迭代器
imges = next(batch_iterator)  #取出最开头的元素

fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
    #在下层中显示生成数据
    plt.subplot(2, 5, i+1)
    plt.imshow(imges[i][0].cpu().detach().numpy(), 'gray')

def Anomaly_score(x, fake_img, z_out_real, D, Lambda=0.1):

   #计算测试图像x与生成图像fake_img在像素层次上的差值的绝对值，并以小批次为单位进行求和计算
    residual_loss = torch.abs(x-fake_img)
    residual_loss = residual_loss.view(residual_loss.size()[0], -1)
    residual_loss = torch.sum(residual_loss, dim=1)

    # 将测试图像x和生成图像fake_img输入判别器D中，并取出特征量图

    _, x_feature = D(x, z_out_real)
    _, G_feature = D(fake_img, z_out_real)

    # 计算测试图像x与生成图像fake_img的特征量的差的绝对值，并以小批次为单位进行求和计算
    discrimination_loss = torch.abs(x_feature-G_feature)
    discrimination_loss = discrimination_loss.view(
        discrimination_loss.size()[0], -1)
    discrimination_loss = torch.sum(discrimination_loss, dim=1)

   #将每个小批次中的两种损失相加
    loss_each = (1-Lambda)*residual_loss + Lambda*discrimination_loss

    #对所有批次中的损失进行计算
    total_loss = torch.sum(loss_each)

    return total_loss, loss_each, residual_loss

#需要检测异常的图像
x = imges[0:5]
x = x.to(device)

#对监督数据的图像进行编码，转换成z，再用生成器G生成图像
z_out_real = E_update(imges.to(device))
imges_reconstract = G_update(z_out_real)

#计算损失值
loss, loss_each, residual_loss_each = Anomaly_score(
    x, imges_reconstract, z_out_real, D_update, Lambda=0.1)

#计算损失值，损失总和
loss_each = loss_each.cpu().detach().numpy()
print("total loss：", np.round(loss_each, 0))

#图像的可视化
fig = plt.figure(figsize=(15, 6))
for i in range(0, 5):
    #在上层中显示训练数据
    plt.subplot(2, 5, i+1)
    plt.imshow(imges[i][0].cpu().detach().numpy(), 'gray')

   #在下层中显示生成数据
    plt.subplot(2, 5, 5+i+1)
    plt.imshow(imges_reconstract[i][0].cpu().detach().numpy(), 'gray')