使用PyTorch构建卷积GAN源码（详细步骤讲解+注释版） 02人脸图片生成上

news2024/11/18 13:56:08

阅读提示：本篇文章的代码为在普通GAN代码上实现人脸图片生成的修改，文章内容仅包含修改内容，全部代码讲解需结合下面的文章阅读。
相关资料链接为：使用PyTorch构建GAN生成对抗
本次训练代码使用了本地GPU计算。

1 CelebADataset类的修改

原则上这一类不需要修改，但为了提升模型运行速度，可以对图片周边适当裁剪，保留五官等重要内容。

# 设置裁剪功能（辅助函数）
def crop_centre(img, new_width, new_height):
    height, width, _ = img.shape
    startx = width//2 - new_width//2
    starty = height//2 - new_height//2
    return img[  starty:starty + new_height, startx:startx + new_width, :]

上面这个函数可以用来从图像的中心裁剪。该函数接收三个参数：

img：原始图像，需要是 numpy 数组形式
new_width：裁剪后图像的新宽度
new_height：裁剪后图像的新高度

该函数通过计算原始图像的中心位置，以及所需裁剪图像的起始位置，从而在 numpy 数组上实现裁剪。最后，函数返回裁剪后的图像。
有了这个函数后，可以在类中预置对图像的裁剪功能，需要对类的__getitem__方法和plot_image方法进行优化。

class CelebADataset(Dataset):

    def __getitem__(self, index):
        if index >= len(self.dataset):
            raise IndexError()
        img = numpy.array(self.dataset[str(index) + '.jpg'])
        img = crop_centre(img, 128, 128)
        return torch.cuda.FloatTensor(img).permute(2,0,1).view(1,3,128,128) / 255.0

    def plot_image(self, index):
        img = numpy.array(self.dataset[str(index)+'.jpg'])
        img = crop_centre(img, 128, 128)
        plt.imshow(img, interpolation='nearest')

2 鉴别器类的修改

鉴别器的网络结构是卷积GAN需要重点修改的地方。此次的卷积GAN设置了3个卷积层和1个全连接层。

class Discriminator(nn.Module):

    def __init__(self):	
    
		self.model = nn.Sequential(
		    
		    nn.Conv2d(3, 256, kernel_size=8, stride=2),
		    nn.BatchNorm2d(256),
		    nn.LeakyReLU(0.2),
		
		    nn.Conv2d(256, 256, kernel_size=8, stride=2),
		    nn.BatchNorm2d(256),
		    nn.LeakyReLU(0.2),
		
		    nn.Conv2d(256, 3, kernel_size=8, stride=2),
		    nn.LeakyReLU(0.2),
		
		    View(3*10*10),
		    nn.Linear(3*10*10, 1),
		    nn.Sigmoid()
	)

经过裁剪的图片的小为128*128；
第一个卷积层使用了256个卷积核，每个卷积核大小为8，步长为2。这一卷积层将会输出256个特征图，特征图的大小为 $\frac{128-8}{2}+1$ ，即61*61；
第二个卷积层使用了256个卷积核，每个卷积核大小为8，步长为2。这一卷积层将会输出256个特征图，特征图的大小为 $\frac{61-8}{2}+1$ ，即27*27；
第二个卷积层使用了3个卷积核，每个卷积核大小为8，步长为2。这一卷积层将会输出3个特征图，特征图的大小为 $\frac{27-8}{2}+1$ ，即10*10；
经过了3层的卷积后，图片的大小已经降到了（3*10*10）。

3 鉴别器测试

修改完鉴别器之后，可以使用真实图像和随即图像，初步判断鉴别器的能力与测试这部分修改后的代码是否存在BUG。

# 鉴别器类建立
D = Discriminator()
D.to(device)

# 测试鉴别器
for image_data_tensor in celeba_dataset:
    # real data
    D.train(image_data_tensor, torch.cuda.FloatTensor([1.0]))
    # fake data
    D.train(generate_random_image((1,3,128,128)), torch.cuda.FloatTensor([0.0]))
    pass

同样，可以查看损失函数的变化情况并使用测试集进行测试。

for image_data_tensor in celeba_dataset:
    # real data
    D.train(image_data_tensor, torch.cuda.FloatTensor([1.0]))
    # fake data
    D.train(generate_random_image((1,3,128,128)), torch.cuda.FloatTensor([0.0]))
    pass

D.plot_progress()

for i in range(4):
  image_data_tensor = celeba_dataset[random.randint(0,20000)]
  print( D.forward( image_data_tensor ).item() )
  pass

for i in range(4):
  print( D.forward( generate_random_image((1,3,128,128))).item() )
  pass