非kaggle内容 一些自学的模型
This commit is contained in:
@@ -0,0 +1,234 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn, optim
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import datasets, transforms
|
||||
from torch.nn import functional as F
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
|
||||
class SinusoidalEmbedding(nn.Module):
|
||||
"""DDPM 标准的正弦位置编码，把时间步 t 编码成向量"""
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, t):
|
||||
half_dim = self.dim // 2
|
||||
emb = math.log(10000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
|
||||
emb = t[:, None] * emb[None, :]
|
||||
return torch.cat([emb.sin(), emb.cos()], dim=-1)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_ch, out_ch, stride=1, time_dim=None):
|
||||
super().__init__()
|
||||
# 第一个卷积：可能改变通道数和步长
|
||||
self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1,
|
||||
padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(out_ch)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(out_ch)
|
||||
|
||||
# 时间嵌入投影：把 t 的嵌入映射到 out_ch，加到特征图上
|
||||
self.time_proj = None
|
||||
if time_dim is not None:
|
||||
self.time_proj = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(time_dim, out_ch),
|
||||
)
|
||||
|
||||
# shortcut：如果维度变化，用1x1卷积（也可加BN）
|
||||
self.shortcut = nn.Sequential()
|
||||
if in_ch != out_ch :
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch)
|
||||
)
|
||||
|
||||
def forward(self, x, t_emb=None):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
# 添加时间嵌入（FiLM 式相加注入）
|
||||
if t_emb is not None and self.time_proj is not None:
|
||||
out = out + self.time_proj(t_emb)[:, :, None, None]
|
||||
|
||||
# 添加 shortcut
|
||||
out += self.shortcut(residual)
|
||||
|
||||
out = self.relu(out) # 标准做法：最后加一次激活
|
||||
return out
|
||||
class DownSample(nn.Module):
|
||||
def __init__(self,ch):
|
||||
super().__init__()
|
||||
self.conv1=nn.Conv2d(ch,ch,2,2)
|
||||
def forward(self,x):
|
||||
return self.conv1(x)
|
||||
|
||||
class UpSample(nn.Module):
|
||||
def __init__(self, ch):
|
||||
super().__init__()
|
||||
# 使用 ch//2 确保整数
|
||||
self.conv1 = nn.ConvTranspose2d(ch, ch//2, 4, 2, 1)
|
||||
def forward(self, x):
|
||||
return self.conv1(x)
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(self, time_dim=64):
|
||||
super().__init__()
|
||||
self.time_emb = SinusoidalEmbedding(time_dim)
|
||||
self.time_mlp = nn.Sequential(
|
||||
nn.Linear(time_dim, time_dim * 4),
|
||||
nn.SiLU(),
|
||||
nn.Linear(time_dim * 4, time_dim * 4),
|
||||
)
|
||||
t_dim = time_dim * 4
|
||||
|
||||
self.conv1 = nn.Conv2d(1, 32, 3, 1, 1)
|
||||
self.res1 = ResidualBlock(32, 64, time_dim=t_dim)
|
||||
self.down1 = DownSample(64)
|
||||
self.res2 = ResidualBlock(64, 128, time_dim=t_dim)
|
||||
self.down2 = DownSample(128)
|
||||
self.res3 = ResidualBlock(128, 256, time_dim=t_dim)
|
||||
self.down3 = DownSample(256)
|
||||
self.res4 = ResidualBlock(256, 512, time_dim=t_dim)
|
||||
|
||||
# --- 上采样部分（重新设计通道匹配）---
|
||||
self.up1 = UpSample(512) # 512 -> 256
|
||||
self.res5 = ResidualBlock(512, 256, time_dim=t_dim) # 拼接后 512 -> 256
|
||||
|
||||
self.up2 = UpSample(256) # 256 -> 128
|
||||
self.res6 = ResidualBlock(256, 128, time_dim=t_dim) # 拼接后 256 -> 128
|
||||
|
||||
self.up3 = UpSample(128) # 128 -> 64
|
||||
self.res7 = ResidualBlock(128, 64, time_dim=t_dim) # 拼接后 128 -> 64
|
||||
|
||||
self.res8 = ResidualBlock(64, 32, time_dim=t_dim)
|
||||
self.conv2 = nn.Conv2d(32, 1, 3, 1, 1)
|
||||
|
||||
def forward(self, x, t):
|
||||
t_emb = self.time_mlp(self.time_emb(t)) # [b, time_dim*4]
|
||||
# ----- 下采样（保存跳跃连接）-----
|
||||
x1 = self.res1(self.conv1(x), t_emb) # [b, 64, 28, 28]
|
||||
x2 = self.res2(self.down1(x1), t_emb) # [b, 128, 14, 14]
|
||||
x3 = self.res3(self.down2(x2), t_emb) # [b, 256, 7, 7]
|
||||
x4 = self.res4(self.down3(x3), t_emb) # [b, 512, 3, 3]
|
||||
|
||||
# ----- 上采样 1（3×3 → 7×7）-----
|
||||
x4_up = self.up1(x4) # [b, 256, 6, 6]
|
||||
x4_up = F.interpolate(x4_up, size=7, mode='bilinear') # [b, 256, 7, 7]
|
||||
x4_cat = torch.cat([x4_up, x3], dim=1) # [b, 512, 7, 7]
|
||||
x3_new = self.res5(x4_cat, t_emb) # [b, 256, 7, 7]
|
||||
|
||||
# ----- 上采样 2（7×7 → 14×14）-----
|
||||
x3_up = self.up2(x3_new) # [b, 128, 14, 14] 尺寸恰好为14
|
||||
x3_cat = torch.cat([x3_up, x2], dim=1) # [b, 256, 14, 14]
|
||||
x2_new = self.res6(x3_cat, t_emb) # [b, 128, 14, 14]
|
||||
|
||||
# ----- 上采样 3（14×14 → 28×28）-----
|
||||
x2_up = self.up3(x2_new) # [b, 64, 28, 28]
|
||||
x2_cat = torch.cat([x2_up, x1], dim=1) # [b, 128, 28, 28]
|
||||
x1_new = self.res7(x2_cat, t_emb) # [b, 64, 28, 28]
|
||||
|
||||
# ----- 最终输出-----
|
||||
x1_new = self.res8(x1_new, t_emb) # [b, 32, 28, 28]
|
||||
out = self.conv2(x1_new) # [b, 1, 28, 28]
|
||||
return out
|
||||
if __name__=='__main__':
|
||||
# ---------- 数据加载 ----------
|
||||
batch_size = 16
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(), # 将 [0,255] 转为 [0,1] 的 Tensor
|
||||
transforms.Normalize((0.5,), (0.5,)), #归一化到[-1,1]
|
||||
]
|
||||
)
|
||||
|
||||
train_dataset = datasets.MNIST(
|
||||
root="./data", train=True, download=True, transform=transform
|
||||
)
|
||||
test_dataset = datasets.MNIST(
|
||||
root="./data", train=False, download=True, transform=transform
|
||||
)
|
||||
|
||||
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
|
||||
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
|
||||
|
||||
# ---------- DDPM 超参数 ----------
|
||||
timesteps = 600
|
||||
betas = torch.linspace(1e-4, 0.02, timesteps) # 线性采样 beta
|
||||
alphas = 1.0 - betas
|
||||
alpha_bar = torch.cumprod(alphas, dim=0) # \bar{alpha}_t
|
||||
|
||||
def q_sample(x_0, t):
|
||||
"""前向扩散：q(x_t | x_0) = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * eps"""
|
||||
alpha_bar_t = alpha_bar[t][:, None, None, None]
|
||||
eps = torch.randn_like(x_0)
|
||||
x_t = torch.sqrt(alpha_bar_t) * x_0 + torch.sqrt(1 - alpha_bar_t) * eps
|
||||
return x_t, eps
|
||||
|
||||
def sample(model, n=16, device="cpu"):
|
||||
"""简易采样：从纯噪声逐步去噪"""
|
||||
model.eval()
|
||||
x = torch.randn(n, 1, 28, 28, device=device)
|
||||
with torch.no_grad():
|
||||
for t in reversed(range(timesteps)):
|
||||
t_batch = torch.full((n,), t, device=device, dtype=torch.long)
|
||||
eps_pred = model(x, t_batch)
|
||||
alpha_t = alphas[t].to(device)
|
||||
alpha_bar_t = alpha_bar[t].to(device)
|
||||
x = (x - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_t)
|
||||
if t > 0:
|
||||
x += torch.sqrt(betas[t].to(device)) * torch.randn_like(x)
|
||||
model.train()
|
||||
return x
|
||||
|
||||
# ---------- 训练 ----------
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model = UNet().to(device)
|
||||
optimizer = optim.Adam(model.parameters(), lr=1e-3)
|
||||
epochs = 20
|
||||
|
||||
for epoch in range(epochs):
|
||||
total_loss = 0.0
|
||||
for x_0, _ in train_loader:
|
||||
x_0 = x_0.to(device)
|
||||
t = torch.randint(0, timesteps, (x_0.size(0),), device=device) # 随机采样时间
|
||||
x_t, eps = q_sample(x_0, t)
|
||||
|
||||
eps_pred = model(x_t, t)
|
||||
loss = F.mse_loss(eps_pred, eps) # 预测噪声，MSE 损失
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += loss.item()
|
||||
|
||||
print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}")
|
||||
|
||||
# ---------- 采样可视化 ----------
|
||||
samples = sample(model, n=16, device=device).cpu().clamp(-1, 1)
|
||||
samples = (samples + 1) / 2 # [-1,1] -> [0,1]
|
||||
|
||||
fig, axes = plt.subplots(4, 4, figsize=(8, 8))
|
||||
for i, ax in enumerate(axes.flat):
|
||||
ax.imshow(samples[i, 0], cmap="gray")
|
||||
ax.axis("off")
|
||||
plt.tight_layout()
|
||||
plt.savefig("samples.png")
|
||||
plt.show()
|
||||
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
import os
|
||||
import argparse
|
||||
import math
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
from torchvision import transforms
|
||||
import numpy as np
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
from PIL import Image
|
||||
|
||||
from main import UNet
|
||||
|
||||
|
||||
def build_schedule(timesteps, device):
|
||||
betas = torch.linspace(1e-4, 0.02, timesteps).to(device)
|
||||
alphas = 1.0 - betas
|
||||
alpha_bar = torch.cumprod(alphas, dim=0)
|
||||
return betas, alphas, alpha_bar
|
||||
|
||||
|
||||
def denoise_step(model, x, t, alphas, alpha_bar, betas):
|
||||
n = x.size(0)
|
||||
t_batch = torch.full((n,), t, device=x.device, dtype=torch.long)
|
||||
eps_pred = model(x, t_batch)
|
||||
alpha_t = alphas[t]
|
||||
alpha_bar_t = alpha_bar[t]
|
||||
x = (x - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_t)
|
||||
if t > 0:
|
||||
x += torch.sqrt(betas[t]) * torch.randn_like(x)
|
||||
return x
|
||||
|
||||
|
||||
def to_grid_image(x, ncols, cmap="gray"):
|
||||
n = x.size(0)
|
||||
nrows = math.ceil(n / ncols)
|
||||
fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 2, nrows * 2))
|
||||
axes = np.atleast_1d(axes).flatten()
|
||||
for i, ax in enumerate(axes):
|
||||
if i < n:
|
||||
ax.imshow(x[i, 0], cmap=cmap)
|
||||
ax.axis("off")
|
||||
plt.subplots_adjust(wspace=0.05, hspace=0.05)
|
||||
fig.canvas.draw()
|
||||
buf = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
|
||||
plt.close(fig)
|
||||
return Image.fromarray(buf)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="DDPM 采样过程逐帧记录")
|
||||
parser.add_argument("--ckpt", type=str, default="model_state_dict.pth")
|
||||
parser.add_argument("--timesteps", type=int, default=600)
|
||||
parser.add_argument("--samples", type=int, default=4)
|
||||
parser.add_argument("--outdir", type=str, default="sampling_frames")
|
||||
parser.add_argument("--gif", action="store_true", help="额外输出 GIF")
|
||||
parser.add_argument("--gif-fps", type=int, default=20)
|
||||
args = parser.parse_args()
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
model = UNet().to(device)
|
||||
state = torch.load(args.ckpt, map_location=device)
|
||||
if isinstance(state, dict) and "state_dict" in state:
|
||||
state = state["state_dict"]
|
||||
model.load_state_dict(state)
|
||||
model.eval()
|
||||
|
||||
betas, alphas, alpha_bar = build_schedule(args.timesteps, device)
|
||||
|
||||
os.makedirs(args.outdir, exist_ok=True)
|
||||
|
||||
x = torch.randn(args.samples, 1, 28, 28, device=device)
|
||||
|
||||
frames = []
|
||||
with torch.no_grad():
|
||||
for t in reversed(range(args.timesteps)):
|
||||
x = denoise_step(model, x, t, alphas, alpha_bar, betas)
|
||||
grid = x.cpu().clamp(-1, 1)
|
||||
grid = (grid + 1) / 2
|
||||
img = to_grid_image(grid, ncols=2)
|
||||
path = os.path.join(args.outdir, f"step_{t:04d}.png")
|
||||
img.save(path)
|
||||
frames.append(path)
|
||||
if t % 50 == 0 or t == 0:
|
||||
print(f"step {t} 已保存 -> {path}")
|
||||
|
||||
print(f"全部 {args.timesteps} 帧已保存到 {args.outdir}/")
|
||||
|
||||
if args.gif:
|
||||
images = [Image.open(p) for p in frames]
|
||||
gif_path = os.path.join(args.outdir, "sampling.gif")
|
||||
images[0].save(
|
||||
gif_path,
|
||||
save_all=True,
|
||||
append_images=images[1:],
|
||||
duration=1000 // args.gif_fps,
|
||||
loop=0,
|
||||
)
|
||||
print(f"GIF 已保存 -> {gif_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 4.1 MiB |
Reference in New Issue
Block a user