生成式AI模型实现MNIST数据增强

提交于 更新于 分类于 ， Disqus：
本文主要是作为我本身的结课的作业的一种提交形式。内容主要包括使用生成式AI模型实现MNIST数据增强。选用CGAN、VAE等生成式AI模型。在深度学习框架PyTorch上基于MNIST训练集优化所选生成式AI模型。基于定量性能指标：[分类准确率]对比无数据增强、有数据增强技术路线性能。并且对所选生成式AI模型分析其数据增强的效果差异并进行对比分析。

文档说明：

​ 本文主要是作为我本身的结课的作业的一种提交形式。内容主要包括使用生成式AI模型实现MNIST数据增强。选用cGAN、VAE等生成式AI模型。在深度学习框架PyTorch上基于MNIST训练集优化所选生成式AI模型。基于定量性能指标：[分类准确率]对比无数据增强、有数据增强技术路线性能。并且对所选生成式AI模型分析其数据增强的效果差异并进行对比分析。

​ 流程图如：

image-20241206142328369

VAE数据增强实现部分：

​ VAE数据增强部分，我这里是分成了两个脚本文件来进行操作，其中一个文件是train_vae.py，这个文件主要是为了训练vae模型。然后还有一个文件，名字是generate_vae.py，该文件主要是利用训练好的vae模型对MNIST数据集进行数据增强，我这里是将MNIST的训练集中的每个样本图片都进行了重参数化（reparameterize）然后输入到Decoder中生成对应的数据样本，同时也是方便我获取生成样本的标签，这是因为VAE模型的Decoder是从潜在空间中生成的样本，本质上是无标签的，所以我要利用原样本的标签和其在潜在空间中的值，这会帮助我生成增强数据集。

VAE训练代码train_vae.py：

以下是我的homeWork/train_vae.py文件的说明：

该脚本在MNIST数据集上训练变分自编码器（VAE）。

依赖项：

  • torch
  • torch.nn
  • torch.optim
  • torch.utils.data
  • torchvision.datasets
  • torchvision.transforms

VAE模型：

VAE模型由编码器和解码器组成。编码器将输入数据压缩到潜在空间表示，解码器从该表示中重构数据。

编码器：

  • 输入：784维向量（展平的28x28图像）
  • 输出：两个20维向量（均值和对数方差）

解码器：

  • 输入：20维潜在向量
  • 输出：784维向量（重构的图像）

重参数化技巧：

为了允许通过随机采样过程进行反向传播，使用重参数化技巧：$ z = \mu + \epsilon \cdot \sigma\ $其中$\epsilon$是从标准正态分布中采样。

损失函数：

VAE的损失函数由重构损失和KL散度组成：

  • 重构损失：衡量重构图像与原始图像的匹配程度。
  • KL散度：衡量潜在空间分布与标准正态分布的接近程度。

训练：

  1. 加载MNIST数据集。
  2. 初始化VAE模型和优化器。
  3. 训练模型若干个epoch，更新模型参数以最小化损失函数。
  4. 将训练好的模型保存到文件。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 定义VAE模型
class VAE(nn.Module):
    def __init__(self, latent_dim=20):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 2 * latent_dim)  # 输出均值和对数方差
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 784),
            nn.Sigmoid()  # 输出值范围在[0, 1]
        )
        self.latent_dim = latent_dim

    def encode(self, x):
        h = self.encoder(x)
        mu, log_var = h.chunk(2, dim=-1)
        return mu, log_var

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x):
        mu, log_var = self.encode(x.view(-1, 784))
        z = self.reparameterize(mu, log_var)
        return self.decode(z), mu, log_var

# 定义损失函数
def vae_loss(recon_x, x, mu, log_var):
    recon_loss = nn.functional.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
    kl_divergence = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return recon_loss + kl_divergence, recon_loss, kl_divergence

if __name__ == '__main__':
    # 加载MNIST数据
    transform = transforms.ToTensor()
    mnist_train = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
    mnist_loader = DataLoader(mnist_train, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)

    # 初始化模型与优化器
    latent_dim = 20
    vae = VAE(latent_dim=latent_dim).to('cuda')
    optimizer = optim.Adam(vae.parameters(), lr=1e-3)

    try:
        vae.load_state_dict(torch.load('vae_gen_mnist.pth', weights_only=True))
        print("Model loaded successfully.")
    except FileNotFoundError:
        print("Model file not found. Initializing model with random weights.")

    # 训练VAE模型
    epochs = 150
    vae.train()
    for epoch in range(epochs):
        train_loss = 0
        for x, _ in mnist_loader:
            x = x.to('cuda')
            optimizer.zero_grad()
            recon_x, mu, log_var = vae(x)
            loss, recon_loss, kl_divergence = vae_loss(recon_x, x, mu, log_var)
            loss.backward()
            train_loss += loss.item()
            optimizer.step()
        print(f"Epoch {epoch + 1}, Total Loss: {train_loss / len(mnist_loader.dataset):.4f}, "
              f"Reconstruction Loss: {recon_loss.item() / len(mnist_loader.dataset):.4f}, "
              f"KL Divergence: {kl_divergence.item() / len(mnist_loader.dataset):.4f}")

    # 保存训练好的模型
    torch.save(vae.state_dict(), 'vae_gen_mnist.pth')

VAE生成代码generate_vae.py:

该脚本使用预训练的变分自编码器（VAE）生成MNIST数据集的样本，并将生成的样本保存为.ubyte文件格式。

依赖项：

  • torch
  • torch.nn
  • numpy
  • matplotlib.pyplot
  • torch.utils.data
  • torchvision.datasets
  • torchvision.transforms
  • torchvision.utils
  • struct

VAE模型：

VAE模型由编码器和解码器组成。编码器将输入数据压缩到潜在空间表示，解码器从该表示中重构数据。

编码器：

  • 输入：784维向量（展平的28x28图像）
  • 输出：两个20维向量（均值和对数方差）

解码器：

  • 输入：20维潜在向量
  • 输出：784维向量（重构的图像）

重参数化技巧：

为了允许通过随机采样过程进行反向传播，使用重参数化技巧：$z = \mu + \epsilon \cdot \sigma$ 其中$\epsilon$是从标准正态分布中采样。

保存增强数据集方法：

save_dataset(images_filepath, labels_filepath, dataset)函数将生成的图像和标签保存为.ubyte文件格式。

  • 参数：
    • images_filepath：图像文件路径
    • labels_filepath：标签文件路径
    • dataset：包含图像和标签的列表

主要步骤：

  1. 定义VAE模型。
  2. 尝试加载预训练（如果有）的VAE模型权重。
  3. 加载MNIST数据集。
  4. 使用VAE生成样本。
  5. 将生成的样本保存为VAE-Generated-images-idx3-ubyte文件和VAE-Generated-labels-idx1-ubyte。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
import struct

# ======================== 定义VAE模型 ========================
class VAE(nn.Module):
    def __init__(self, latent_dim=20):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 2 * latent_dim)  # 输出均值和对数方差
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 784),
            nn.Sigmoid()  # 输出值范围在[0, 1]
        )
        self.latent_dim = latent_dim

    def encode(self, x):
        h = self.encoder(x)
        mu, log_var = h.chunk(2, dim=-1)
        return mu, log_var

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x):
        mu, log_var = self.encode(x.view(-1, 784))
        z = self.reparameterize(mu, log_var)
        return self.decode(z), mu, log_var

# ======================== 保存增强数据集方法 ========================
def save_dataset(images_filepath, labels_filepath, dataset):
    """保存合并后的数据集到 .ubyte 文件"""
    images = []
    labels = []

    for image, label in dataset:
        images.append(image.numpy())
        labels.append(label)

    images = np.stack(images).astype(np.uint8)  # 转为 [N, 28, 28] 的 uint8 数组
    labels = np.array(labels, dtype=np.uint8)  # 转为 uint8 的标签

    print(f"Labels before saving: {labels.shape}")
    print(f"Unique labels before saving: {np.unique(labels)}")

    # 写入图像文件
    with open(images_filepath, 'wb') as f:
        f.write(struct.pack('>IIII', 2051, len(images), 28, 28))  # Magic number 2051 for images
        f.write(images.tobytes())

    # 写入标签文件
    with open(labels_filepath, 'wb') as f:
        f.write(struct.pack('>II', 2049, len(labels)))  # Magic number 2049 for labels
        f.write(labels.tobytes())

if __name__ == '__main__':
    # 加载训练好的模型
    latent_dim = 20
    vae = VAE(latent_dim=latent_dim).to('cuda')
    vae.load_state_dict(torch.load('vae_gen_mnist.pth', weights_only=True))
    vae.eval()

    # 加载MNIST数据集
    transform = transforms.ToTensor()
    mnist_train = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
    mnist_loader = DataLoader(mnist_train, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)

    # ======================== 使用VAE生成样本 ========================
    generated_samples = []
    generated_labels = []
    with torch.no_grad():
        for x, y in mnist_loader:
            x = x.to('cuda')
            mu, log_var = vae.encode(x.view(-1, 784))
            z = vae.reparameterize(mu, log_var)
            generated_images = vae.decode(z).cpu().view(-1, 28, 28).numpy()  # 转为 28x28 的 NumPy 数组
            generated_samples.extend(generated_images)
            generated_labels.extend(y.numpy())

    # 将列表转换为单一 numpy 数组
    generated_samples = np.array(generated_samples)  # 转为 (N, 28, 28) 的 numpy 数组
    generated_labels = np.array(generated_labels, dtype=np.uint8)  # 转为 uint8 的标签数组

    # 确保生成的样本数据在 [0, 255] 范围内
    generated_samples = (generated_samples * 255).astype(np.uint8)

    # 保存生成的样本到 .ubyte 文件
    generated_images_path = './enhanced_mnist/VAE-Generated-images-idx3-ubyte'
    generated_labels_path = './enhanced_mnist/VAE-Generated-labels-idx1-ubyte'

    # 将生成的样本和标签打包成数据集
    generated_dataset = [(torch.tensor(image, dtype=torch.uint8), label) for image, label in zip(generated_samples, generated_labels)]

    # 使用 save_dataset 函数保存生成的数据集
    save_dataset(generated_images_path, generated_labels_path, generated_dataset)

    print(f"生成的数据集已保存到文件：")
    print(f" - 图像文件: {generated_images_path}")
    print(f" - 标签文件: {generated_labels_path}")

CGAN数据增强实现部分：

CGAN训练代码train_gan.py:

依赖项：

torch
torchvision

CGAN模型：

生成器设计（Generator）:

​ 生成器模型用于生成与真实图像相似的图像。其结构如下：

嵌入层：
  • label_emb：将类别标签嵌入到与噪声向量相同的维度中。
  • 参数：num_classes（类别数），num_classes（嵌入向量维度）。
全连接层：
  • fc：将噪声向量和嵌入标签连接起来，并通过全连接层进行处理。
  • 全连接层参数：输入维度 input_size + num_classes，输出维度 num_feature。
卷积层：
  • conv1_g：包含卷积层、批归一化层和ReLU激活函数。
  • 卷积层参数：输入通道数 1，输出通道数 50，卷积核大小 3，填充 1。
  • conv2_g：包含卷积层、批归一化层和ReLU激活函数。
  • 卷积层参数：输入通道数 50，输出通道数 25，卷积核大小 3，填充 1。
  • conv3_g：包含卷积层和Tanh激活函数。
  • 卷积层参数：输入通道数 25，输出通道数 1，卷积核大小 2，步幅 2。

判别器设计（Discriminator）：

​ 判别器模型用于区分真实图像和生成图像。其结构如下：

嵌入层：
  • label_emb：将类别标签嵌入到与图像大小相同的向量中。
  • 参数：num_classes（类别数），28 * 28（嵌入向量维度）。
卷积层：
  • conv1：包含卷积层、LeakyReLU激活函数和最大池化层。
  • 卷积层参数：输入通道数 2，输出通道数 32，卷积核大小 5，填充 2。
  • conv2：包含卷积层、LeakyReLU激活函数和最大池化层。
  • 卷积层参数：输入通道数 32，输出通道数 64，卷积核大小 5，填充 2。
全连接层：
  • fc：包含全连接层、LeakyReLU激活函数和Sigmoid激活函数。
  • 全连接层参数：输入维度 64 * 7 * 7，输出维度 1024 和 1。

损失函数：

GAN模型的损失函数是二元交叉熵损失函数（Binary Cross-Entropy Loss），定义如下：

1
criterion = nn.BCELoss()

在训练过程中，判别器和生成器的损失分别计算如下：

  • 判别器损失（d_loss）：

    1
    2
    3
    d_loss_real = criterion(real_out, real_label)
    d_loss_fake = criterion(fake_out, fake_label)
    d_loss = d_loss_real + d_loss_fake

  • 生成器损失（g_loss）：

    1
    g_loss = criterion(output, real_label)

训练：

  1. 加载MNIST数据集。
  2. 初始化GAN模型和优化器。
  3. 训练模型若干个epoch，更新模型参数以最小化损失函数。
  4. 将训练好的模型保存到文件。

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.autograd import Variable
import os

# # ======================== 创建保存图像的文件夹 ========================
# if not os.path.exists('./GAN'):
#     os.mkdir('./GAN')


# ======================== 定义图像转换函数 ========================
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 = 150  # 训练epoch
z_dimension = 100  # 噪声维度
num_classes = 10  # 类别数

# ======================== 图像处理 ========================
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# ======================== 加载MNIST数据集 ========================
mnist = datasets.MNIST(
    root='./data', train=True, transform=transform, download=True)
dataloader = torch.utils.data.DataLoader(
    dataset=mnist, batch_size=batch_size, shuffle=True)


# ======================== 定义判别器模型 ========================
class Discriminator(nn.Module):
    def __init__(self, num_classes):
        super(Discriminator, self).__init__()
        self.label_emb = nn.Embedding(num_classes, 28 * 28)
        self.conv1 = nn.Sequential(
            nn.Conv2d(2, 32, 5, padding=2),  # batch, 32, 28, 28
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, stride=2)  # batch, 32, 14, 14
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 5, padding=2),  # batch, 64, 14, 14
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, stride=2)  # batch, 64, 7, 7
        )
        self.fc = nn.Sequential(
            nn.Linear(64 * 7 * 7, 1024),
            nn.LeakyReLU(0.2),
            nn.Linear(1024, 1),
            nn.Sigmoid()
        )

    def forward(self, img, labels):
        label_embedding = self.label_emb(labels).view(labels.size(0), 1, 28, 28)
        x = torch.cat((img, label_embedding), 1)
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


# ======================== 定义生成器模型 ========================
class Generator(nn.Module):
    def __init__(self, input_size, num_feature, num_classes):
        super(Generator, self).__init__()
        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.fc = nn.Linear(input_size + num_classes, num_feature)  # batch, 3136=1x56x56
        self.br = nn.Sequential(
            nn.BatchNorm2d(1),
            nn.ReLU(True)
        )
        self.conv1_g = nn.Sequential(
            nn.Conv2d(1, 50, 3, stride=1, padding=1),  # batch, 50, 56, 56
            nn.BatchNorm2d(50),
            nn.ReLU(True)
        )
        self.conv2_g = nn.Sequential(
            nn.Conv2d(50, 25, 3, stride=1, padding=1),  # batch, 25, 56, 56
            nn.BatchNorm2d(25),
            nn.ReLU(True)
        )
        self.conv3_g = nn.Sequential(
            nn.Conv2d(25, 1, 2, stride=2),  # batch, 1, 28, 28
            nn.Tanh()
        )

    def forward(self, noise, labels):
        label_embedding = self.label_emb(labels)
        x = torch.cat((noise, label_embedding), -1)
        x = self.fc(x)
        x = x.view(x.size(0), 1, 56, 56)
        x = self.br(x)
        x = self.conv1_g(x)
        x = self.conv2_g(x)
        x = self.conv3_g(x)
        return x

if __name__ == '__main__':
    # ======================== 初始化模型 ========================
    discriminator = Discriminator(num_classes)
    generator = Generator(z_dimension, 3136,num_classes)
    if torch.cuda.is_available():
        discriminator = discriminator.cuda()
        generator = generator.cuda()

    # ======================== 定义损失函数和优化器 ========================
    # 使用二元交叉熵损失函数
    criterion = nn.BCELoss()
    d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=0.0001)
    g_optimizer = torch.optim.Adam(generator.parameters(), lr=0.0001)

    try:
        generator.load_state_dict(torch.load('./GAN_generator.pth', weights_only=True))
        discriminator.load_state_dict(torch.load('./GAN_discriminator.pth', weights_only=True))
        print("\n--------Model restored--------\n")
    except FileNotFoundError:
        print("\n--------Model not restored: File not found--------\n")
    except Exception as e:
        print(f"\n--------Model not restored: {e}--------\n")
    # ======================== 训练模型 ========================
    for epoch in range(num_epoch):
        for i, (img, labels) in enumerate(dataloader):
            num_img = img.size(0)
            real_img = Variable(img).cuda()
            real_label = Variable(torch.ones(num_img)).cuda().unsqueeze(1)
            fake_label = Variable(torch.zeros(num_img)).cuda().unsqueeze(1)
            labels = Variable(labels).cuda()

            # Train Discriminator
            real_out = discriminator(real_img, labels)
            d_loss_real = criterion(real_out, real_label)
            real_scores = real_out  # Closer to 1 means better

            z = Variable(torch.randn(num_img, z_dimension)).cuda()
            fake_img = generator(z, labels)
            fake_out = discriminator(fake_img, labels)
            d_loss_fake = criterion(fake_out, fake_label)
            fake_scores = fake_out  # Closer to 0 means better

            d_loss = d_loss_real + d_loss_fake
            d_optimizer.zero_grad()
            d_loss.backward()
            d_optimizer.step()

            # Train Generator
            z = Variable(torch.randn(num_img, z_dimension)).cuda()
            fake_img = generator(z, labels)
            output = discriminator(fake_img, labels)
            g_loss = criterion(output, real_label)

            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()))
    # ======================== 保存模型 ========================
    torch.save(generator.state_dict(), './GAN_generator.pth')
    torch.save(discriminator.state_dict(), './GAN_discriminator.pth')

CGAN生成代码generator_GAN.py：

依赖项

该项目需要以下依赖项：

1
2
3
torch
torchvision
numpy

CGAN模型

​ 条件生成对抗网络 (CGAN) 是一种在训练过程中也利用标签的 GAN。生成器 - 给定标签和随机数组作为输入，该网络生成与对应相同标签的训练数据观察具有相同结构的数据。生成器用于生成与真实图像相似的图像，判别器用于区分真实图像和生成图像。在本文件中因为已经默认使用了MNIST数据集对模型进行过了训练，所以本文件中就没有用到辨别器（Discriminator），只用到了生成器。

生成器模型的结构如下：

  1. 嵌入层：

    • label_emb：将类别标签嵌入到与噪声向量相同的维度中。
    • 参数：num_classes（类别数），num_classes（嵌入向量维度）。

  2. 全连接层：

    • fc：将噪声向量和嵌入标签连接起来，并通过全连接层进行处理。
    • 参数：输入维度 input_size + num_classes，输出维度 num_feature。

  3. 卷积层：

    • conv1_g：包含卷积层、批归一化层和ReLU激活函数。
      • 参数：输入通道数 1，输出通道数 50，卷积核大小 3，填充 1。
    • conv2_g：包含卷积层、批归一化层和ReLU激活函数。
      • 参数：输入通道数 50，输出通道数 25，卷积核大小 3，填充 1。
    • conv3_g：包含卷积层和Tanh激活函数。
      • 参数：输入通道数 25，输出通道数 1，卷积核大小 2，步幅 2。

保存数据集的方法

该项目提供了一个方法来保存生成的数据集到 .ubyte 文件中：

1
2
3
4
5
6
7
8
def save_dataset(images_filepath, labels_filepath, images, labels):
    with open(images_filepath, 'wb') as f:
        f.write(struct.pack('>IIII', 2051, len(images), 28, 28))  # Magic number 2051 for images
        f.write(images.tobytes())

    with open(labels_filepath, 'wb') as f:
        f.write(struct.pack('>II', 2049, len(labels)))  # Magic number 2049 for labels
        f.write(labels.tobytes())

主要步骤

  1. 导入必要的库：

    • 导入 torch、torch.nn、numpy 等库。

  2. 定义生成器模型：

    • Generator 类包含嵌入层、全连接层和卷积层，用于生成图像。

  3. 设置超参数：

    • 定义噪声维度、类别数、生成样本数量和输出目录。

  4. 加载生成器模型：

    • 初始化生成器模型，并加载预训练的模型权重。

  5. 生成数据集：

    • 使用生成器模型生成指定数量的样本和标签。

  6. 保存生成的样本到 .ubyte 文件：

    • 使用 save_dataset 方法将生成的样本和标签保存到文件GAN-Generated-images-idx3-ubyte和GAN-Generated-labels-idx1-ubyte中。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import torch
from torch.autograd import Variable
import os
import numpy as np
import struct
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, input_size, num_feature, num_classes):
        super(Generator, self).__init__()
        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.fc = nn.Linear(input_size + num_classes, num_feature)  # batch, 3136=1x56x56
        self.br = nn.Sequential(
            nn.BatchNorm2d(1),
            nn.ReLU(True)
        )
        self.conv1_g = nn.Sequential(
            nn.Conv2d(1, 50, 3, stride=1, padding=1),  # batch, 50, 56, 56
            nn.BatchNorm2d(50),
            nn.ReLU(True)
        )
        self.conv2_g = nn.Sequential(
            nn.Conv2d(50, 25, 3, stride=1, padding=1),  # batch, 25, 56, 56
            nn.BatchNorm2d(25),
            nn.ReLU(True)
        )
        self.conv3_g = nn.Sequential(
            nn.Conv2d(25, 1, 2, stride=2),  # batch, 1, 28, 28
            nn.Tanh()
        )

    def forward(self, noise, labels):
        label_embedding = self.label_emb(labels)
        x = torch.cat((noise, label_embedding), -1)
        x = self.fc(x)
        x = x.view(x.size(0), 1, 56, 56)
        x = self.br(x)
        x = self.conv1_g(x)
        x = self.conv2_g(x)
        x = self.conv3_g(x)
        return x

# 设置超参数
z_dimension = 100  # 噪声维度
num_classes = 10  # 类别数
num_samples = 60000  # 生成样本数量
output_dir = './enhanced_mnist'  # 输出目录

# 创建输出目录
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 加载生成器模型
generator = Generator(z_dimension, 3136, num_classes)
generator.load_state_dict(torch.load('./GAN_generator.pth', weights_only=True))
generator.eval()

if torch.cuda.is_available():
    generator = generator.cuda()

# 生成数据集
generated_samples = []
generated_labels = []
for i in range(num_samples):
    z = Variable(torch.randn(1, z_dimension)).cuda()
    labels = Variable(torch.randint(0, num_classes, (1,))).cuda()
    fake_img = generator(z, labels)
    fake_img = fake_img.cpu().data.numpy().squeeze() * 255  # 转为 numpy 数组并缩放到 [0, 255]
    generated_samples.append(fake_img.astype(np.uint8))
    generated_labels.append(labels.cpu().item())

# 将生成的样本和标签打包成数据集
generated_samples = np.stack(generated_samples)
generated_labels = np.array(generated_labels, dtype=np.uint8)

# 保存生成的样本到 .ubyte 文件
def save_dataset(images_filepath, labels_filepath, images, labels):
    with open(images_filepath, 'wb') as f:
        f.write(struct.pack('>IIII', 2051, len(images), 28, 28))  # Magic number 2051 for images
        f.write(images.tobytes())

    with open(labels_filepath, 'wb') as f:
        f.write(struct.pack('>II', 2049, len(labels)))  # Magic number 2049 for labels
        f.write(labels.tobytes())

generated_images_path = os.path.join(output_dir, 'GAN-Generated-images-idx3-ubyte')
generated_labels_path = os.path.join(output_dir, 'GAN-Generated-labels-idx1-ubyte')

save_dataset(generated_images_path, generated_labels_path, generated_samples, generated_labels)

print(f"生成的数据集已保存到文件：")
print(f" - 图像文件: {generated_images_path}")
print(f" - 标签文件: {generated_labels_path}")

三种数据集可视化对比展示：

原始MNIST数据集展示图：

​ plot_2024-12-12 00-02-25_0

VAE增强数据集展示图：

plot_2024-12-12 23-34-41_2

CGAN增强数据集展示图：

plot_2024-12-12 23-34-41_1

分别用三种数据集训练分类模型：

分类模型选择：LeNet5

我的这里的深度学习的分类模型选择的是LeNet5模型，当然LeNet5模型本身可以对图像的横向纵向进行特征提取就可以达到99%的正确识别率具体为：

LeNet5 模型

LeNet5 是一个经典的卷积神经网络（CNN）模型，主要用于图像分类任务。其结构如下：

  1. 卷积层 1 (conv1)：

    • 输入通道数：1
    • 输出通道数：6
    • 卷积核大小：5

  2. 卷积层 2 (conv2)：

    • 输入通道数：6
    • 输出通道数：16
    • 卷积核大小：5

  3. 全连接层 1 (fc1)：

    • 输入维度：16 * 4 * 4
    • 输出维度：120

  4. 全连接层 2 (fc2)：

    • 输入维度：120
    • 输出维度：84

  5. 全连接层 3 (fc3)：

    • 输入维度：84
    • 输出维度：10

前向传播过程

  1. 输入图像通过第一个卷积层 (conv1)，然后经过 ReLU 激活函数。
  2. 经过最大池化层 (max_pool2d)。
  3. 通过第二个卷积层 (conv2)，然后经过 ReLU 激活函数。
  4. 再次经过最大池化层 (max_pool2d)。
  5. 将特征图展平为一维向量。
  6. 通过第一个全连接层 (fc1)，然后经过 ReLU 激活函数。
  7. 通过第二个全连接层 (fc2)，然后经过 ReLU 激活函数。
  8. 通过第三个全连接层 (fc3)，输出分类结果。

类代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, 16 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

损失函数和优化器：

  1. 设置损失函数：

    • criterion = nn.CrossEntropyLoss()：使用交叉熵损失函数（Cross-Entropy Loss），这是一个常用的分类任务损失函数。

  2. 设置优化器：

    • optimizer = optim.Adam(leNet.parameters(), lr=learning_rate)：使用Adam优化器，并设置学习率。Adam优化器是一种自适应学习率优化算法，适用于大多数深度学习模型。
1
2
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(leNet.parameters(), lr=learning_rate)

代码：

三种训练LeNet5的代码都一样只有数据集加载处有些改动，写的时候为了方便调试拆成了3个单独的文件，其中一个利用GAN训练增强数据集GAN_Train_LeNet_Combined.py的代码如下，（另外两个VAE和Original都类似，仅仅只有加载数据集过程中有些小区别）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, ConcatDataset
import numpy as np

# ======================== 定义LeNet模型 ========================
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, 16 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# ======================== 加载增强的数据集 ========================
def load_enhanced_dataset(images_path, labels_path):
    with open(images_path, 'rb') as f:
        images = np.frombuffer(f.read(), dtype=np.uint8, offset=16).reshape(-1, 28, 28)
    with open(labels_path, 'rb') as f:
        labels = np.frombuffer(f.read(), dtype=np.uint8, offset=8)
    return [(torch.tensor(image, dtype=torch.float32).unsqueeze(0), label) for image, label in zip(images, labels)]

if __name__ == '__main__':
    # ======================== 设置超参数 ========================
    batch_size = 64
    learning_rate = 0.001
    num_epochs = 50

    # ======================== 数据预处理 ========================
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # ======================== 加载MNIST数据集 ========================
    train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
    # test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)

    # ======================== 加载增强的数据集 ========================
    enhanced_images_path = './enhanced_mnist/GAN-Generated-images-idx3-ubyte'
    enhanced_labels_path = './enhanced_mnist/GAN-Generated-labels-idx1-ubyte'
    enhanced_dataset = load_enhanced_dataset(enhanced_images_path, enhanced_labels_path)

    # ======================== 合并数据集 ========================
    combined_train_dataset = ConcatDataset([train_dataset, enhanced_dataset])
    train_loader = DataLoader(combined_train_dataset, batch_size=batch_size, shuffle=True)
    # test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # ======================== 初始化模型、损失函数和优化器 ========================
    leNet = LeNet5().to('cuda')
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(leNet.parameters(), lr=learning_rate)

    try:
        leNet.load_state_dict(torch.load('GAN_leNet_mnist.pth', weights_only=True))
        print("Model loaded successfully.")
    except FileNotFoundError:
        print("Model file not found. Initializing model with random weights.")

    # ======================== 训练模型 ========================
    for epoch in range(num_epochs):
        leNet.train()
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to('cuda'), labels.to('cuda')

            optimizer.zero_grad()
            outputs = leNet(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")

    # ======================== 保存模型 ========================
    torch.save(leNet.state_dict(), 'GAN_leNet_mnist.pth')

效果对比

这里基于MNIST测试集对比三种训练集得到的模型做分类准确率评估评估效果如下：

plot_2024-12-12 23-34-41_3

三种数据集分别训练LeNet5模型进行分类的成功率如下：

原始数据集：98.98%

利用VAE增强的数据集：98.80%

利用GAN增强的数据集：98.37%

结果分析

​ 利用增强数据集训练的 LeNet 模型，分类效果略微下降了一些，但整体分类成功率仍能保持在 98% 以上。这可能是因为：

  1. VAE 和 CGAN 是基于训练集中的数据进行样本生成的，其生成的样本往往会使整个样本集的数据特征更加集中于某些特定方面。当生成的数据过于充分时，这种特性在训练 LeNet-5 时可能进一步削弱其泛化能力，从而导致过拟合，在测试集上表现为分类能力下降和性能减退。这一问题在 LeNet-5 这样的复杂多层卷积网络中尤为显著。
  2. 我认为，部分过拟合现象可能也与数据集本身过于简单有关。以 MNIST 数据集为例，其本身较为简单，而且训练集已经包含了 60,000 条数据，数量相当充足。在数据增强后，训练集的样本数量更是达到了 120,000 条。假设这些增强数据均为有效样本，那么对于像 LeNet-5 这样专注于分类任务的卷积神经网络来说，更容易出现过拟合现象。
  3. CGAN 的工作模式可能会加重过拟合现象。这是因为 GAN 模式中的 Discriminator 是基于训练集对 Generator 生成的数据进行辨别的，而这种对抗性机制会导致 Discriminator 和 Generator 在训练集上表现得过于“精准”，从而更容易出现过拟合问题。