基于U-Net的diffusion model还原图像

代码有参考，忘记链接了，找到后会贴出。

/一些说明/

U-Net连接部分选用双线性插值：

crop1 = F.interpolate(enc3, size=dec4.shape[2:], mode='bilinear', align_corners=True)

数据集选用CIFA10。

输入图像处理部分：transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))数值来源模型论文。

预处理模型选用vgg16。

loss函数选用MSE损失与感知损失。

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torchvision import models
from torchvision.models import VGG16_Weights
from PIL import Image
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# U-Net结构
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        # 编码器部分
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
            nn.ReLU()
        )

        # 解码器部分
        self.deconv1 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
        self.relu1 = nn.ReLU()
        self.deconv2 = nn.ConvTranspose2d(512, 128, kernel_size=4, stride=2, padding=1)  # 输入512（256+256）
        self.relu2 = nn.ReLU()
        self.deconv3 = nn.ConvTranspose2d(25