PyTorch实现自编码器

以下是一个使用PyTorch实现自编码器的示例代码,该代码包括三个自编码器和一些辅助函数,用于训练和测试自编码器。

案例1

  1. import torch
  2. import torch.nn as nn
  3. import torch.optim as optim
  4. import torchvision.transforms as transforms
  5. import torchvision.datasets as datasets
  6. import numpy as np
  7. import matplotlib.pyplot as plt
  8. # Define the Stacked Autoencoder class
  9. class StackedAutoencoder(nn.Module):
  10. def __init__(self, input_dim, hidden_dims):
  11. super(StackedAutoencoder, self).__init__()
  12. self.input_dim = input_dim
  13. self.hidden_dims = hidden_dims
  14. # Define the encoder layers
  15. self.encoder1 = nn.Linear(input_dim, hidden_dims[0])
  16. self.encoder2 = nn.Linear(hidden_dims[0], hidden_dims[1])
  17. self.encoder3 = nn.Linear(hidden_dims[1], hidden_dims[2])
  18. # Define the decoder layers
  19. self.decoder3 = nn.Linear(hidden_dims[2], hidden_dims[1])
  20. self.decoder2 = nn.Linear(hidden_dims[1], hidden_dims[0])
  21. self.decoder1 = nn.Linear(hidden_dims[0], input_dim)
  22. # Define the activation function
  23. self.activation = nn.ReLU()
  24. def encoder(self, x):
  25. z1 = self.activation(self.encoder1(x))
  26. z2 = self.activation(self.encoder2(z1))
  27. z3 = self.activation(self.encoder3(z2))
  28. return z3
  29. def decoder(self, z):
  30. xhat3 = self.activation(self.decoder3(z))
  31. xhat2 = self.activation(self.decoder2(xhat3))
  32. xhat1 = self.decoder1(xhat2)
  33. return xhat1
  34. def forward(self, x):
  35. z = self.encoder(x)
  36. xhat = self.decoder(z)
  37. return xhat
  38. def get_encoder_output(self, x):
  39. return self.encoder(x)
  40. # Define the training function
  41. def train(model, train_loader, num_epochs, learning_rate):
  42. # Define the loss function and optimizer
  43. criterion = nn.MSELoss()
  44. optimizer = optim.Adam(model.parameters(), lr=learning_rate)
  45. # Train the model
  46. for epoch in range(num_epochs):
  47. for data in train_loader:
  48. # Get the input data and target data
  49. inputs, targets = data
  50. inputs, targets = inputs.view(-1, 28*28), targets.view(-1, 28*28)
  51. # Zero the gradients
  52. optimizer.zero_grad()
  53. # Forward pass
  54. outputs = model(inputs)
  55. loss = criterion(outputs, targets)
  56. # Backward pass and optimization
  57. loss.backward()
  58. optimizer.step()
  59. # Print the loss after each epoch
  60. print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
  61. # Define the test function
  62. def test(model, test_loader):
  63. # Define the loss function
  64. criterion = nn.MSELoss()
  65. # Evaluate the model
  66. test_loss = 0
  67. with torch.no_grad():
  68. for data in test_loader:
  69. # Get the input data and target data
  70. inputs, targets = data
  71. inputs, targets = inputs.view(-1, 28*28), targets.view(-1, 28*28)
  72. # Forward pass
  73. outputs = model(inputs)
  74. test_loss += criterion(outputs, targets).item()
  75. # Print the average test loss
  76. test_loss /= len(test_loader.dataset)
  77. print('Average Test Loss: {:.4f}'.format(test_loss))

主程序

  1. # Define the main function
  2. def main():
  3. # Set the device
  4. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  5. # Set the hyperparameters
  6. input_dim = 28*28
  7. hidden_dims = [256, 128, 64]
  8. num_epochs = 10
  9. batch_size = 128
  10. learning_rate = 0.001
  11. # Download the MNIST dataset and create data loaders
  12. train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
  13. test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)
  14. train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
  15. test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
  16. # Create the Stacked Autoencoder model and move it to the device
  17. model = StackedAutoencoder(input_dim, hidden_dims).to(device)
  18. # Train the model
  19. train(model, train_loader, num_epochs, learning_rate)
  20. # Test the model
  21. test(model, test_loader)
  22. # Generate a random image and its reconstruction
  23. with torch.no_grad():
  24. z = torch.randn(1, hidden_dims[-1]).to(device)
  25. xhat = model.decoder(z)
  26. xhat = xhat.view(28, 28).cpu().numpy()
  27. plt.imshow(xhat, cmap='gray')
  28. plt.show()
  29. if __name__ == '__main__':
  30. main()

main() 函数中,首先设置了设备,然后定义了超参数,接着下载 MNIST 数据集并创建数据加载器。然后创建了自编码器模型,并将其移动到设备上。接下来调用 train() 函数进行训练,然后调用 test() 函数进行测试。最后生成一个随机图像并进行重构,然后显示出来。

train() 函数中,定义了损失函数和优化器,然后对模型进行了训练。在 test() 函数中,定义了损失函数,并对模型进行了测试。

test() 函数中,定义了损失函数,并对模型进行了测试。测试过程与训练过程类似,但是不需要进行梯度更新。最后返回测试损失的平均值。


案例2

  1. import torch
  2. import torch.nn as nn
  3. import torch.optim as optim
  4. import torchvision
  5. import torchvision.transforms as transforms
  6. # 数据预处理
  7. transform = transforms.Compose([
  8. transforms.ToTensor(),
  9. transforms.Normalize((0.5,), (0.5,))
  10. ])
  11. # 加载MNIST数据集
  12. trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
  13. trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
  14. # 定义自编码器模型
  15. class StackedAutoencoder(nn.Module):
  16. def __init__(self, input_dim, hidden_dims):
  17. super(StackedAutoencoder, self).__init__()
  18. self.encoder = nn.Sequential(
  19. nn.Linear(input_dim, hidden_dims[0]),
  20. nn.ReLU()
  21. )
  22. self.decoder = nn.Sequential(
  23. nn.Linear(hidden_dims[0], input_dim),
  24. nn.ReLU()
  25. )
  26. def forward(self, x):
  27. x = self.encoder(x)
  28. x = self.decoder(x)
  29. return x
  30. # 训练自编码器
  31. input_dim = 28 * 28 # MNIST图像尺寸为28x28
  32. hidden_dims = [256, 128] # 隐藏层维度
  33. model = StackedAutoencoder(input_dim, hidden_dims)
  34. criterion = nn.MSELoss()
  35. optimizer = optim.Adam(model.parameters(), lr=0.001)
  36. num_epochs = 10
  37. for epoch in range(num_epochs):
  38. running_loss = 0.0
  39. for images, _ in trainloader:
  40. images = images.view(images.size(0), -1)
  41. optimizer.zero_grad()
  42. outputs = model(images)
  43. loss = criterion(outputs, images)
  44. loss.backward()
  45. optimizer.step()
  46. running_loss += loss.item()
  47. epoch_loss = running_loss / len(trainloader)
  48. print(f"Epoch {
  49. epoch+1}/{
  50. num_epochs}, Loss: {
  51. epoch_loss:.4f}")
  52. print("Training finished!")
  53. # 使用自编码器进行图像重建
  54. testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
  55. testloader = torch.utils.data.DataLoader(testset, batch_size=10, shuffle=False)
  56. dataiter = iter(testloader)
  57. images, labels = dataiter.next()
  58. images = images.view(images.size(0), -1)
  59. outputs = model(images)
  60. # 可视化原始图像和重建图像
  61. import matplotlib.pyplot as plt
  62. def imshow(img):
  63. img = img / 2 + 0.5 # 反归一化
  64. npimg = img.numpy()
  65. plt.imshow(np.transpose(npimg, (1, 2, 0)))
  66. plt.axis('off')
  67. plt.show()
  68. # 显示原始图像
  69. imshow(torchvision.utils.make_grid(images.view(-1, 1, 28, 28)))
  70. # 显示重建图像
  71. imshow(torchvision.utils.make_grid(outputs.view(-1, 1, 28, 28)))

这个案例实现了一个简单的两层自编码器,用于对MNIST数据集中的手写数字图像进行重建。首先定义了一个StackedAutoencoder类,它包含一个编码器和一个解码器,其中编码器是一个包含ReLU激活函数的全连接层序列,解码器也是一个包含ReLU激活函数的全连接层序列。然后,通过使用MSE损失函数和Adam优化器对模型进行训练。在训练过程中,将图像展平为784维的向量,并将其输入到模型中,然后计算重建图像与原始图像之间的损失,并进行反向传播和参数更新。最后,使用训练好的模型对一批测试图像进行重建,并将原始图像和重建图像可视化。

自编码器可以包含更多的隐藏层,更复杂的结构和更多的训练步骤,以更好地学习数据的表示。此外,还可以尝试在编码器和解码器之间添加dropout层、使用不同的激活函数等来改进模型的性能。

发表评论

表情:
评论列表 (有 0 条评论,48人围观)

还没有评论,来说两句吧...

相关阅读

    相关 PyTorch实现联邦学习堆叠编码器

    联邦学习是一种用于训练分布在不同设备或地点的模型的技术,其中数据分布在不同的设备上,且不会离开设备。每个设备只训练其本地数据的模型,并将更新的模型参数传递给服务器,服务器对这些

    相关 编码器

    无监督学习介绍 Markdown Code 内容简介 目前许多有监督学习算法,如 SVM,DNN 或是 boosting,决策树等,都在工业界分类或决策任务上取得了