非kaggle内容 一些自学的模型

Mnist-DDPM/record_sampling.py

@@ -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()