使用 CUDA 在 PyTorch 中实现卷积自动编码器

263 阅读4分钟

自动编码器是一种神经网络架构,用于无监督学习任务,例如数据压缩、降维和数据去噪。该架构由两个主要组件组成:编码器和解码器。网络的编码器部分将输入数据压缩为低维表示,而网络的解码器部分根据该低维表示重建原始输入数据。

卷积自动编码器 (CAE) 是一种自动编码器,是一种深度学习神经网络架构,通常用于无监督学习任务,例如图像压缩和去噪。它是传统自动编码器架构的扩展,将卷积层合并到网络的编码器和解码器部分。

与自动编码器相同,卷积自动编码器架构也由两个主要组件组成:编码器和解码器。网络的编码器部分使用卷积层和池化操作处理输入图像,以生成图像的低维特征表示。网络的解码器部分采用这种低维特征表示,并使用反卷积层将其上采样回原始输入图像大小。网络的最终输出是尽可能接近原始输入图像的重建图像。

卷积自动编码器的训练过程与传统自动编码器的训练过程类似。网络经过训练,可使用均方误差 (MSE) 或二元交叉熵 (BCE) 等损失函数来最小化原始输入图像和重建输出图像之间的差异。一旦训练完成,网络的编码器部分可用于特征提取,网络的解码器部分可用于图像生成或重建。

卷积自动编码器在各种计算机视觉任务中显示了令人印象深刻的结果,包括图像压缩、去噪和特征提取。它们还被用于图像检索、对象识别和异常检测等各种应用中。

Pytorch 中的实现:

算法

  1.  使用 PyTorch 的 ImageFolder 类加载数据集并定义数据加载器。
  2.  通过创建包含编码器和解码器的自动编码器类来定义卷积自动编码器架构,每个编码器和解码器都具有卷积层和池化层。
  3. 初始化自动编码器模型,并使用 to() 方法将其移动到 GPU(如果可用)。
  4. 定义训练期间使用的损失函数和优化器。通常,使用均方误差 (MSE) 损失,Adam 优化器是深度学习任务的热门选择。
  5. 设置要训练的纪元数并开始训练循环。
  6. 在每个 epoch 中,迭代数据加载器的批次,将数据移动到 GPU,并执行前向传播以获得自动编码器的输出。
  7. 使用损失函数计算输出和输入之间的损失。
  8. 执行反向传播以计算模型参数相对于损失的梯度。\
  9.  使用优化器根据计算的梯度更新模型参数。
  10. 在每个时期后打印损失以监控训练进度。
  11. 使用 state_dict() 方法将训练后的模型保存到文件中。

代码:

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 定义自动编码器架构
class Autoencoder(nn.Module):
	def __init__(self):
		super(Autoencoder, self).__init__()
		self.encoder = nn.Sequential(
			nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(kernel_size=2, stride=2),
			nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(kernel_size=2, stride=2)
		)
		self.decoder = nn.Sequential(
			nn.ConvTranspose2d(8, 16,
							kernel_size=3,
							stride=2,
							padding=1,
							output_padding=1),
			nn.ReLU(),
			nn.ConvTranspose2d(16, 3,
							kernel_size=3,
							stride=2,
							padding=1,
							output_padding=1),
			nn.Sigmoid()
		)
		
	def forward(self, x):
		x = self.encoder(x)
		x = self.decoder(x)
		return x

初始化自动编码器

model = Autoencoder()

定义变换

transform = transforms.Compose([
	transforms.Resize((64, 64)),
	transforms.ToTensor(),
])

加载数据集

train_dataset = datasets.Flowers102(root='flowers',
									split='train',
									transform=transform,
									download=True)
test_dataset = datasets.Flowers102(root='flowers',
								split='test',
								transform=transform)

定义数据加载器

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
										batch_size=128,
										shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
										batch_size=128)

将模型移动到GPU

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
model.to(device)

定义损失函数和优化器

criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

训练自动编码器

num_epochs = 50
for epoch in range(num_epochs):
	for data in train_loader:
		img, _ = data
		img = img.to(device)
		optimizer.zero_grad()
		output = model(img)
		loss = criterion(output, img)
		loss.backward()
		optimizer.step()
	if epoch % 5== 0:
		print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

保存模型

torch.save(model.state_dict(), 'conv_autoencoder.pth')

输出

cuda
Epoch [1/50], Loss: 0.0919
Epoch [6/50], Loss: 0.0746
Epoch [11/50], Loss: 0.0362
Epoch [16/50], Loss: 0.0239
Epoch [21/50], Loss: 0.0178
Epoch [26/50], Loss: 0.0154
Epoch [31/50], Loss: 0.0144
Epoch [36/50], Loss: 0.0124
Epoch [41/50], Loss: 0.0127
Epoch [46/50], Loss: 0.0101

用解码图像绘制原始图像

with torch.no_grad():
	for data, _ in test_loader:
		data = data.to(device)
		recon = model(data)
		break
		
import matplotlib.pyplot as plt
plt.figure(dpi=250)
fig, ax = plt.subplots(2, 7, figsize=(15, 4))
for i in range(7):
	ax[0, i].imshow(data[i].cpu().numpy().transpose((1, 2, 0)))
	ax[1, i].imshow(recon[i].cpu().numpy().transpose((1, 2, 0)))
	ax[0, i].axis('OFF')
	ax[1, i].axis('OFF')
plt.show()

输出

image.png