结合PyTorch、Diffusion Model和Matplotlib利用mnist训练集进行大量训练以达到生成逼真的手写数字，要求训练模型可视化并可可视化损失随时间的变化曲线同时，该代码还将手写数字样本保存为图像文件，方便查看

由于任务较为复杂，我们将代码分成几部分来讲解。

## 1. 数据预处理

首先，我们需要下载 MNIST 数据集并进行预处理。

import torch
from torchvision import datasets, transforms

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 下载训练集
train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transform)

# 加载训练集
train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True)

## 2. 构建模型

接下来，我们需要构建 Diffusion Model。

import torch.nn as nn

class DiffusionModel(nn.Module):
    def __init__(self, input_size=784, hidden_size=256, output_size=784, num_layers=2):
        super(DiffusionModel, self).__init__()
        
        # 编码器
        encoder_layers = []
        for i in range(num_layers):
            if i == 0:
                encoder_layers.append(nn.Linear(input_size, hidden_size))
            else:
                encoder_layers.append(nn.Linear(hidden_size, hidden_size))
            encoder_layers.append(nn.ReLU())
        self.encoder = nn.Sequential(*encoder_layers)
        
        # 解码器
        decoder_layers = []
        for i in range(num_layers):
            if i == 0:
                decoder_layers.append(nn.Linear(hidden_size, output_size))
            else:
                decoder_layers.append(nn.Linear(hidden_size, hidden_size))
            decoder_layers.append(nn.ReLU())
        decoder_layers.append(nn.Linear(hidden_size, input_size))
        decoder_layers.append(nn.Tanh())
        self.decoder = nn.Sequential(*decoder_layers)
        
    def forward(self, x):
        z = self.encoder(x)
        y = self.decoder(z)
        return y

## 3. 训练模型

现在，我们可以开始训练模型了。

import time
import matplotlib.pyplot as plt

# 定义模型和优化器
model = DiffusionModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# 训练模型
num_epochs = 50
losses = []
start_time = time.time()
for epoch in range(num_epochs):
    for batch_idx, (data, _) in enumerate(train_loader):
        data = data.view(data.shape[0], -1)
        optimizer.zero_grad()
        recon_data = model(data)
        loss = criterion(recon_data, data)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

# 打印训练时间
end_time = time.time()
print('Training Time: {:.2f}s'.format(end_time - start_time))

# 可视化损失函数
plt.plot(losses)
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.show()

## 4. 保存模型和样本

最后，我们可以保存模型和生成手写数字样本。

import os
import torchvision.utils as vutils

# 创建目录
os.makedirs('images', exist_ok=True)

# 保存模型
torch.save(model.state_dict(), 'diffusion_model.pth')

# 生成样本
num_samples = 64
z = torch.randn(num_samples, 256)
samples = model.decoder(z)
samples = samples.view(num_samples, 1, 28, 28)
vutils.save_image(samples, 'images/samples.png', normalize=True, nrow=8)

现在，你可以在 images/samples.png 中查看生成的手写数字样本了。