{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":95504,"databundleVersionId":11357675,"sourceType":"competition"},{"sourceId":311174,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":263909,"modelId":285006}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# MNIST 生成任务\n\n在这个部分，你们会尝试运行扩散模型，并完成一个简单的生成任务，这个生成过程与AI绘画的原理类似。 你所需要的代码已经被提供，你只需要完成代码的填充和运行。","metadata":{}},{"cell_type":"markdown","source":"导入必需的库","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torchvision import datasets, transforms\nfrom torch.utils.data import DataLoader, TensorDataset\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch.nn.functional as F","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:05.694658Z","iopub.execute_input":"2025-03-31T11:49:05.694920Z","iopub.status.idle":"2025-03-31T11:49:11.295027Z","shell.execute_reply.started":"2025-03-31T11:49:05.694897Z","shell.execute_reply":"2025-03-31T11:49:11.294114Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"指定运行设备","metadata":{}},{"cell_type":"code","source":"# 指定在GPU上运行程序\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.295866Z","iopub.execute_input":"2025-03-31T11:49:11.296152Z","iopub.status.idle":"2025-03-31T11:49:11.345028Z","shell.execute_reply.started":"2025-03-31T11:49:11.296134Z","shell.execute_reply":"2025-03-31T11:49:11.344097Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"构建网络模型","metadata":{}},{"cell_type":"code","source":"class Denoiser(nn.Module):\n    def __init__(self,\n                 img_channels=1,            \n                 label_dim=10,              \n                 base_channels=64,       \n                 sigma_data=0.5):           \n        super().__init__()\n        self.img_channels = img_channels\n        self.label_dim = label_dim\n        self.sigma_data = sigma_data\n        self.noise_embed = nn.Sequential(\n            nn.Linear(1, 128),            \n            nn.SiLU(),\n            nn.Linear(128, 256),           \n            nn.SiLU(),\n            nn.Linear(256, 256),            \n            nn.SiLU(),\n            nn.Linear(256, 256)            \n        )\n        \n\n        self.label_embed = nn.Sequential(\n            nn.Embedding(label_dim, 128),   \n            nn.SiLU(),\n            nn.Linear(128, 256),          \n            nn.SiLU(),\n            nn.Linear(256, 256)            \n        )\n        \n   \n        self.res1 = ResBlock(img_channels, base_channels)\n        self.res2 = ResBlock(base_channels, base_channels*2)\n        self.res3 = ResBlock(base_channels*2, base_channels*4)  \n        self.res4 = ResBlock(base_channels*4, base_channels*4)   \n        self.res5 = ResBlock(base_channels*4, base_channels*2)   \n        self.res6 = ResBlock(base_channels*2, base_channels)\n        self.cond_proj1 = nn.Linear(256, base_channels)\n        self.cond_proj2 = nn.Linear(256, base_channels*2)\n        self.cond_proj3 = nn.Linear(256, base_channels*4)\n        self.cond_proj4 = nn.Linear(256, base_channels*4)\n        self.cond_proj5 = nn.Linear(256, base_channels*2)\n        self.cond_proj6 = nn.Linear(256, base_channels)\n        self.output = nn.Conv2d(base_channels, img_channels, kernel_size=3, padding=1)\n    \n    def forward(self, x, sigma, class_labels):\n        x = x.to(torch.float32)\n        sigma = sigma.to(torch.float32).reshape(-1, 1)\n        class_labels = class_labels.long()\n        \n        c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)\n        c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt()\n        c_in = 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt()\n        c_noise = sigma.log() / 4\n        noise_emb = self.noise_embed(c_noise)\n        label_emb = self.label_embed(class_labels)\n        emb = noise_emb + label_emb\n        h = c_in.view(-1, 1, 1, 1) * x\n        \n        h = self.res1(h, self.cond_proj1(emb))\n        h = self.res2(h, self.cond_proj2(emb))\n        h = self.res3(h, self.cond_proj3(emb))\n        h = self.res4(h, self.cond_proj4(emb))  \n        h = self.res5(h, self.cond_proj5(emb)) \n        h = self.res6(h, self.cond_proj6(emb))\n        \n        h = self.output(h)\n        \n        # 输出处理\n        D_x = c_skip.view(-1, 1, 1, 1) * x + c_out.view(-1, 1, 1, 1) * h\n        \n        return D_x\n        \nclass ResBlock(nn.Module):\n    def __init__(self, in_ch, out_ch):\n        super().__init__()\n        self.main = nn.Sequential(\n            nn.GroupNorm(min(8, in_ch), in_ch),\n            nn.SiLU(),\n            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n            nn.GroupNorm(min(8, out_ch), out_ch),\n            nn.SiLU(),\n            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n            # 增加一个卷积层\n            nn.GroupNorm(min(8, out_ch), out_ch),\n            nn.SiLU(),\n            nn.Conv2d(out_ch, out_ch, 3, padding=1)\n        )\n        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()\n        \n    def forward(self, x, cond):\n        h = self.main(x)\n        h = h + cond.view(-1, cond.shape[1], 1, 1)\n        return h + self.skip(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.346185Z","iopub.execute_input":"2025-03-31T11:49:11.346548Z","iopub.status.idle":"2025-03-31T11:49:11.364800Z","shell.execute_reply.started":"2025-03-31T11:49:11.346513Z","shell.execute_reply":"2025-03-31T11:49:11.364016Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"定义损失函数","metadata":{}},{"cell_type":"code","source":"class EDMLoss:\n    def __init__(self, P_mean=-1.2, P_std=1.2, sigma_data=0.5):\n        self.P_mean = P_mean\n        self.P_std = P_std\n        self.sigma_data = sigma_data\n\n    def __call__(self, net, images, labels=None, augment_pipe=None):\n        rnd_normal = torch.randn([images.shape[0]], device=images.device)\n        sigma = (rnd_normal * self.P_std + self.P_mean).exp()\n        weight = (sigma ** 2 + self.sigma_data ** 2) / (sigma * self.sigma_data) ** 2\n        y = images\n        n = torch.randn_like(y) * sigma.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)\n        D_yn = net(y + n, sigma, labels.to(torch.float32))\n        weight = weight.view(-1, 1, 1, 1)\n        loss = weight * ((D_yn - y) ** 2)\n        return loss.mean()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.365484Z","iopub.execute_input":"2025-03-31T11:49:11.365720Z","iopub.status.idle":"2025-03-31T11:49:11.378797Z","shell.execute_reply.started":"2025-03-31T11:49:11.365694Z","shell.execute_reply":"2025-03-31T11:49:11.378028Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 让我们了解一下扩散模型的原理。 \n\n通常，对于我们想要采样的分布，其具体的数学形式往往是未知的，我们只能得到一些从该分布中采样的样本。 为了从未知分布中采样，我们往往需要学习数据的分布，并从学习到的分布中采样。 \n\n训练去噪器的目的是获得分布函数score function,分布的对数密度函数的梯度,它是这个扩散过程重要的组成部分\n\n$score(x,\\sigma) = \\nabla_{x}\\log p(x,\\sigma)$\n\n$dx =-\\dot{\\sigma}(t)\\sigma(t)\\nabla_{x}\\log p(x,\\sigma)dt$\n\nscore function的输入是带噪声的图像与噪声的水平, 返回图像空间上的一个向量,这个方向从含噪分布指向原始数据分布,利用它我们可以从初始噪声回到数据分布.","metadata":{}},{"cell_type":"markdown","source":"定义生成过程","metadata":{}},{"cell_type":"code","source":"# 输出采样图像\ndef conditional_sample(model, #去噪器\n                       latents, #初始噪声\n                       labels, #图像类别\n                       num_steps, #采样步数\n                       sigma_min, #最小噪声水平\n                       sigma_max, #最大噪声水平\n                       rho, #噪声水平调度的非线性程度\n                       alpha):\n    model.eval()\n    # 初始噪声\n    x = latents\n    # 生成过程\n\n    # 生成幂律分布的时间节点\n    step_indices = torch.arange(num_steps, dtype=torch.float32, device=device)\n    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho\n    t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])])  # t_N = 0\n\n    # 每一个时间节点 进行一次迭代\n    for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):\n        x_cur = x\n        h = t_next - t_cur\n        t_cur_batch = t_cur.repeat(latents.shape[0])\n        labels_batch = labels.repeat(latents.shape[0])\n        # 计算当前时间节点 网络基于目前的图像与噪声水平的输出\n        denoised = model(x_cur, t_cur_batch, labels_batch).to(torch.float32)\n        # 基于网络的输出 指出一个图像空间中指向数据分布的方向\n        d_cur = (x_cur - denoised) / t_cur\n        # 向这个方向进一步优化图像\n        x_prime = x_cur + alpha * h * d_cur\n        t_prime = t_cur + alpha * h\n        t_prime_batch = t_prime.repeat(latents.shape[0])\n        if t_next != 0:\n            # 计算校正后的输出\n            denoised_prime = model(x_prime, t_prime_batch, labels_batch).to(torch.float32)\n            d_prime = (x_prime - denoised_prime) / t_prime\n            # 基于校正后的输出 进一步优化图像\n            x = x_cur + h * ((1 - 1 / (2 * alpha)) * d_cur + 1 / (2 * alpha) * d_prime)\n        # 在最后一步禁用校正\n        else:\n            x = x_cur + h * d_cur\n    \n    return x.cpu().squeeze()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.379616Z","iopub.execute_input":"2025-03-31T11:49:11.379881Z","iopub.status.idle":"2025-03-31T11:49:11.398989Z","shell.execute_reply.started":"2025-03-31T11:49:11.379854Z","shell.execute_reply":"2025-03-31T11:49:11.398200Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"数据加载与预处理","metadata":{}},{"cell_type":"code","source":"# 读取数据集\n\ndef load_MNIST(shuffle=True, batch_size=256):\n    transform = transforms.Compose([\n        transforms.ToTensor(),\n        transforms.Resize((28, 28)),\n        transforms.Normalize((0.5,), (0.5,))\n    ])\n    \n    # 加载数据集\n    train_dataset = datasets.MNIST(\n        root='data_download/mnist',\n        train=True,\n        transform=transform,\n        download=True\n    )\n    \n    # 创建DataLoader\n    train_loader = DataLoader(\n        dataset=train_dataset,\n        batch_size=batch_size,\n        shuffle=shuffle,\n        num_workers=4,\n        pin_memory=False,\n        drop_last=False\n    )\n    \n    return train_loader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.400922Z","iopub.execute_input":"2025-03-31T11:49:11.401177Z","iopub.status.idle":"2025-03-31T11:49:11.417473Z","shell.execute_reply.started":"2025-03-31T11:49:11.401157Z","shell.execute_reply":"2025-03-31T11:49:11.416663Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"训练及加载网络参数","metadata":{}},{"cell_type":"code","source":"# 实例化一个网络类\ndef train_model(learning_rate, num_epochs, batch_size):\n    model = Denoiser().to(device)\n    # 打印这个网络的参数数量\n    total_params = sum(p.numel() for p in model.parameters())\n    print(f\"Total parameters: {total_params}\")\n    # 指定网络的optimizer和scheduler\n    optimizer = optim.AdamW(model.parameters(), lr=learning_rate)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)\n    # 指定loss类型\n    loss_fn = EDMLoss()\n    # 选择是否进行训练\n    train = False\n    # 选择是否读取网络参数\n    load_model = True\n    if load_model:\n        # 读取参数\n        model.load_state_dict(torch.load('/kaggle/input/denoiser-model/pytorch/default/1/kaggle_model.pth', \n                                         map_location=torch.device('cpu'),\n                                         weights_only = True))\n    if train:\n        # 进行训练\n        train_loader = load_MNIST(batch_size=batch_size)\n        for epoch in range(num_epochs):\n            total_loss = 0\n            for x, labels in train_loader:\n                x = x.to(device)\n                labels = labels.to(device)\n                \n                optimizer.zero_grad()\n                loss = loss_fn(model, x, labels)\n                loss.backward()\n                optimizer.step()\n                \n                total_loss += loss.item()\n            scheduler.step()\n            print(f\"Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}\")\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.418368Z","iopub.execute_input":"2025-03-31T11:49:11.418553Z","iopub.status.idle":"2025-03-31T11:49:11.433435Z","shell.execute_reply.started":"2025-03-31T11:49:11.418536Z","shell.execute_reply":"2025-03-31T11:49:11.432646Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 任务\n\n现在，我们提供给你网络以及训练的代码，不过训练一个去噪器会消耗大量时间，我们不建议在课堂上多次尝试，所以这里提供一个已经训练完成的模型参数。\n\n你可以考虑在已经训练好的模型上进一步训练，这也许会有效的改善生成样本的质量。 但除此之外，你需要调整采样过程的参数，这同样对生成的样本有很大的影响。\n\n你需要找到合适的参数，生成有效的图像。\n\n我们希望你可以利用之前学到的知识，找到可能有问题的地方，自由发挥进行修改，当你能够对任意的类别生成一张合适的图像时，你就可以进入第二部分的测试了。","metadata":{}},{"cell_type":"markdown","source":"## 训练参数设定","metadata":{}},{"cell_type":"code","source":"# 如果你不准备进行训练 可以不用修改这一部分参数\nbatch_size = 256 # 一次梯度下降的batch大小\nlearning_rate = 1e-3 # 初始学习率大小\nnum_epochs = 2 # 总共进行100个epoch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.434076Z","iopub.execute_input":"2025-03-31T11:49:11.434283Z","iopub.status.idle":"2025-03-31T11:49:11.452442Z","shell.execute_reply.started":"2025-03-31T11:49:11.434264Z","shell.execute_reply":"2025-03-31T11:49:11.451790Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 采样参数设定","metadata":{}},{"cell_type":"code","source":"# 这一部分是采样过程的参数设定\nnum_steps = 5 # 采样过程进行的步数\nsigma_min = 1 # 采样过程的输出的噪声水平 不能设置为零\nsigma_max = 2 # 初始噪声的规模 \nrho = 7 # 噪声水平调度的非线性程度\nalpha = 1 # Heun积分控制参数","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.453257Z","iopub.execute_input":"2025-03-31T11:49:11.453572Z","iopub.status.idle":"2025-03-31T11:49:11.466824Z","shell.execute_reply.started":"2025-03-31T11:49:11.453527Z","shell.execute_reply":"2025-03-31T11:49:11.466131Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"生成并绘制图像","metadata":{}},{"cell_type":"code","source":"def plot_mnist_images(images, nrows=2, ncols=5, title=None):\n    # 创建大图\n    fig = plt.figure(figsize=(ncols*2.5, nrows*2.5), dpi=150)\n    \n    # 遍历图像\n    for i, image_tensor in enumerate(images):\n        image_tensor = image_tensor.unsqueeze(0)\n        if i >= nrows*ncols:  # 超出设定的行列数，停止绘制\n            break\n            \n        # 创建子图\n        plt.subplot(nrows, ncols, i+1)\n        if isinstance(image_tensor, torch.Tensor):\n            if image_tensor.is_cuda:\n                image_tensor = image_tensor.cpu()\n            image_tensor = (image_tensor + 1) / 2\n            image_tensor = torch.clamp(image_tensor, 0, 1)\n            if image_tensor.shape[0] == 1:\n                img = image_tensor.squeeze(0).numpy()\n            else:\n                img = image_tensor.permute(1, 2, 0).numpy()\n        else:\n            img = (image_tensor + 1) / 2\n            img = np.clip(img, 0, 1)\n            \n        # 绘制图像\n        plt.imshow(img, cmap='gray', interpolation='nearest')\n        plt.title(f\"sample {i+1}\", fontsize=10)\n        plt.axis('off')\n    \n    if title:\n        plt.suptitle(title, fontsize=16)\n        \n    plt.subplots_adjust(hspace=0.3, wspace=0.1)\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.467515Z","iopub.execute_input":"2025-03-31T11:49:11.467793Z","iopub.status.idle":"2025-03-31T11:49:11.487218Z","shell.execute_reply.started":"2025-03-31T11:49:11.467765Z","shell.execute_reply":"2025-03-31T11:49:11.486630Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"latents = torch.randn(10, 1, 28, 28, device=device) * sigma_max\n# 指定你想要生成的数字","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 实例化一个去噪器并读取参数\nmodel = train_model(learning_rate, num_epochs, batch_size)\n# 生成一个sigma_max水平的初始噪声\nlables = torch.tensor([9], device=device)\n# 进行生成\nwith torch.no_grad():\n    images = conditional_sample(model, latents, lables, num_steps, sigma_min, sigma_max, rho, alpha)\nplot_mnist_images(images)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-31T11:49:11.487939Z","iopub.execute_input":"2025-03-31T11:49:11.488144Z","iopub.status.idle":"2025-03-31T11:49:13.865285Z","shell.execute_reply.started":"2025-03-31T11:49:11.488125Z","shell.execute_reply":"2025-03-31T11:49:13.864366Z"}},"outputs":[],"execution_count":null}]}