以下是一个使用
PyTorch
实现堆叠自编码器的示例代码,该代码包括三个自编码器和一些辅助函数,用于训练和测试堆叠自编码器。
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import numpy as np
import matplotlib.pyplot as plt
# Define the Stacked Autoencoder class
class StackedAutoencoder(nn.Module):
def __init__(self, input_dim, hidden_dims):
super(StackedAutoencoder, self).__init__()
self.input_dim = input_dim
self.hidden_dims = hidden_dims
# Define the encoder layers
self.encoder1 = nn.Linear(input_dim, hidden_dims[0])
self.encoder2 = nn.Linear(hidden_dims[0], hidden_dims[1])
self.encoder3 = nn.Linear(hidden_dims[1], hidden_dims[2])
# Define the decoder layers
self.decoder3 = nn.Linear(hidden_dims[2], hidden_dims[1])
self.decoder2 = nn.Linear(hidden_dims[1], hidden_dims[0])
self.decoder1 = nn.Linear(hidden_dims[0], input_dim)
# Define the activation function
self.activation = nn.ReLU()
def encoder(self, x):
z1 = self.activation(self.encoder1(x))
z2 = self.activation(self.encoder2(z1))
z3 = self.activation(self.encoder3(z2))
return z3
def decoder(self, z):
xhat3 = self.activation(self.decoder3(z))
xhat2 = self.activation(self.decoder2(xhat3))
xhat1 = self.decoder1(xhat2)
return xhat1
def forward(self, x):
z = self.encoder(x)
xhat = self.decoder(z)
return xhat
def get_encoder_output(self, x):
return self.encoder(x)
# Define the training function
def train(model, train_loader, num_epochs, learning_rate):
# Define the loss function and optimizer
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# Train the model
for epoch in range(num_epochs):
for data in train_loader:
# Get the input data and target data
inputs, targets = data
inputs, targets = inputs.view(-1, 28*28), targets.view(-1, 28*28)
# Zero the gradients
optimizer.zero_grad()
# Forward pass
outputs = model(inputs)
loss = criterion(outputs, targets)
# Backward pass and optimization
loss.backward()
optimizer.step()
# Print the loss after each epoch
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
# Define the test function
def test(model, test_loader):
# Define the loss function
criterion = nn.MSELoss()
# Evaluate the model
test_loss = 0
with torch.no_grad():
for data in test_loader:
# Get the input data and target data
inputs, targets = data
inputs, targets = inputs.view(-1, 28*28), targets.view(-1, 28*28)
# Forward pass
outputs = model(inputs)
test_loss += criterion(outputs, targets).item()
# Print the average test loss
test_loss /= len(test_loader.dataset)
print('Average Test Loss: {:.4f}'.format(test_loss))
主程序
# Define the main function
def main():
# Set the device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Set the hyperparameters
input_dim = 28*28
hidden_dims = [256, 128, 64]
num_epochs = 10
batch_size = 128
learning_rate = 0.001
# Download the MNIST dataset and create data loaders
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# Create the Stacked Autoencoder model and move it to the device
model = StackedAutoencoder(input_dim, hidden_dims).to(device)
# Train the model
train(model, train_loader, num_epochs, learning_rate)
# Test the model
test(model, test_loader)
# Generate a random image and its reconstruction
with torch.no_grad():
z = torch.randn(1, hidden_dims[-1]).to(device)
xhat = model.decoder(z)
xhat = xhat.view(28, 28).cpu().numpy()
plt.imshow(xhat, cmap='gray')
plt.show()
if __name__ == '__main__':
main()
在
main()
函数中,首先设置了设备,然后定义了超参数,接着下载MNIST
数据集并创建数据加载器。然后创建了堆叠自编码器模型,并将其移动到设备上。接下来调用train()
函数进行训练,然后调用test()
函数进行测试。最后生成一个随机图像并进行重构,然后显示出来。
在
train()
函数中,定义了损失函数和优化器,然后对模型进行了训练。在test()
函数中,定义了损失函数,并对模型进行了测试。
在
test()
函数中,定义了损失函数,并对模型进行了测试。测试过程与训练过程类似,但是不需要进行梯度更新。最后返回测试损失的平均值。