自动编码器代码
时间: 2023-07-09 20:04:58 浏览: 40
下面是一个简单的自动编码器的代码示例,使用PyTorch框架实现:
```
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.Linear(28*28, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 12),
nn.ReLU(),
nn.Linear(12, 3)
)
self.decoder = nn.Sequential(
nn.Linear(3, 12),
nn.ReLU(),
nn.Linear(12, 64),
nn.ReLU(),
nn.Linear(64, 128),
nn.ReLU(),
nn.Linear(128, 28*28),
nn.Sigmoid()
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
# 定义超参数
batch_size = 32
learning_rate = 1e-3
num_epochs = 10
# 定义数据加载器和损失函数
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
model = Autoencoder()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# 训练自动编码器模型
for epoch in range(num_epochs):
for data in train_loader:
img, _ = data
img = img.view(img.size(0), -1)
output = model(img)
loss = criterion(output, img)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
# 测试自动编码器模型
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
with torch.no_grad():
for data in test_loader:
img, _ = data
img = img.view(img.size(0), -1)
output = model(img)
loss = criterion(output, img)
print('Test Loss: {:.4f}'.format(loss.item()))
```
这个示例中定义了一个包含三个隐层的自动编码器模型,使用MNIST数据集进行训练和测试。在训练过程中使用了均方误差损失函数和Adam优化器。在测试过程中计算了重构误差。