稀疏自动编码器pytorch代码

是使用PyTorch实现稀疏自编码器的代码：

import torch
import torch.nn as nn
import torch.optim as optim

class SparseAutoencoder(nn.Module):
    def __init__(self, input_size, hidden_size, sparsity_target):
        super(SparseAutoencoder, self).__init__()
        self.encoder = nn.Linear(input_size, hidden_size)
        self.decoder = nn.Linear(hidden_size, input_size)
        self.sparsity_target = sparsity_target
        self.sparse_loss = nn.KLDivLoss(reduction='sum')

    def forward(self, x):
        hidden = torch.sigmoid(self.encoder(x))
        output = torch.sigmoid(self.decoder(hidden))
        return output, hidden

    def training_step(self, x, optimizer):
        optimizer.zero_grad()
        output, hidden = self.forward(x)
        reconstruction_loss = nn.MSELoss()(output, x)
        sparsity_loss = self.sparse_loss(torch.tensor([self.sparsity_target]), torch.mean(hidden, dim=0))
        loss = reconstruction_loss + sparsity_loss
        loss.backward()
        optimizer.step()
        return loss.item()

# 定义模型参数
input_size = 784
hidden_size = 128
sparsity_target = 0.1
learning_rate = 0.01
num_epochs = 10

# 加载数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

# 初始化模型和优化器
model = SparseAutoencoder(input_size, hidden_size, sparsity_target)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    total_loss = 0
    for batch in train_loader:
        images, _ = batch
        images = images.view(images.size(0), -1)
        loss = model.training_step(images, optimizer)
        total_loss += loss
    print("Epoch:", epoch+1, "Loss:", total_loss/len(train_loader))

# 测试模型
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

with torch.no_grad():
    for batch in test_loader:
        images, _ = batch
        images = images.view(images.size(0), -1)
        outputs, _ = model(images)
        reconstruction_loss = nn.MSELoss()(outputs, images)
        print("Reconstruction Loss:", reconstruction_loss.item())