Hugging Face's logo

Buckets:
hf-doc-build
/
doc-dev

Files
xet
hf-doc-build/doc-dev / diffusers /pr_14355 /zh /_app /immutable /nodes /52.CYEzuQ0U.js
HuggingFaceDocBuilder's picture
HuggingFaceDocBuilder
download
raw
55.9 kB
import"../chunks/DsnmJJEf.js";import{i as X,h as S,C as W,H as a,a as l,E as Y,s as z}from"../chunks/DdZvggmf.js";import{p as v,o as O,s,f as D,a as G,b as x,c as f,n as H}from"../chunks/BbekZcyp.js";import{D as q}from"../chunks/Bje6zuSL.js";const K='{"title":"训练扩散模型","local":"训练扩散模型","sections":[{"title":"训练配置","local":"训练配置","sections":[],"depth":2},{"title":"加载数据集","local":"加载数据集","sections":[],"depth":2},{"title":"创建 UNet2DModel","local":"创建-unet2dmodel","sections":[],"depth":2},{"title":"创建调度器","local":"创建调度器","sections":[],"depth":2},{"title":"训练模型","local":"训练模型","sections":[],"depth":2},{"title":"下一步","local":"下一步","sections":[],"depth":2}],"depth":1}';var L=f('<meta name="hf:doc:metadata"/>'),P=f('<p></p> <!> <!> <!> <p>无条件图像生成是扩散模型最常见的应用之一，它会生成与训练数据集风格相似的图像。通常来说，在某个特定数据集上微调预训练模型能得到最好的结果。你可以在 <a href="https://huggingface.co/search/full-text?q=unconditional-image-generation&amp;type=model" rel="nofollow">Hub</a> 上找到很多现成检查点；如果找不到满意的，也完全可以自己训练一个！</p> <p>这篇教程会教你如何在 <a href="https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset" rel="nofollow">Smithsonian Butterflies</a> 数据集的一个子集上，从零开始训练一个 <code>UNet2DModel</code>，生成属于你自己的 🦋 蝴蝶图像 🦋。</p> <blockquote class="tip"><p>💡 这篇训练教程基于 <a href="https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb" rel="nofollow">Training with 🧨 Diffusers</a> notebook 编写。如果你想了解更多背景，例如扩散模型的工作原理，也推荐一起看看这个 notebook。</p></blockquote> <p>开始之前，请确认已经安装了 🤗 Datasets，用来加载和预处理图像数据集；以及 🤗 Accelerate，用来简化任意数量 GPU 上的训练。下面这条命令也会安装 <a href="https://www.tensorflow.org/tensorboard" rel="nofollow">TensorBoard</a> 来可视化训练指标（你也可以使用 <a href="https://docs.wandb.ai/" rel="nofollow">Weights & Biases</a> 跟踪训练）。</p> <!> <p>我们也很鼓励你把模型分享给社区。为此，你需要登录自己的 Hugging Face 账号（如果还没有，可以在 <a href="https://hf.co/join" rel="nofollow">这里</a> 创建）。你可以在 notebook 中登录，系统会提示你输入 token。请确保这个 token 具有写入权限。</p> <!> <p>或者在终端里登录：</p> <!> <p>由于模型检查点通常比较大，建议安装 <a href="https://git-lfs.com/" rel="nofollow">Git-LFS</a> 来管理这些大文件：</p> <!> <!> <p>为了方便起见，我们先创建一个 <code>TrainingConfig</code> 类，把训练超参数放在一起（你可以按需调整）：</p> <!> <!> <p>你可以很轻松地通过 🤗 Datasets 加载 <a href="https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset" rel="nofollow">Smithsonian Butterflies</a> 数据集：</p> <!> <blockquote class="tip"><p>💡 你也可以从 <a href="https://huggingface.co/huggan" rel="nofollow">HugGan Community Event</a> 找到更多数据集，或者通过本地 <a href="https://huggingface.co/docs/datasets/image_dataset#imagefolder" rel="nofollow"><code>ImageFolder</code></a> 使用自己的数据集。如果你使用 HugGan Community Event 里的数据集，把 <code>config.dataset_name</code> 设为对应数据集的 repository id；如果你使用自己的图像，就设为 <code>imagefolder</code>。</p></blockquote> <p>🤗 Datasets 使用 <code>Image</code> 特性自动解码图像数据，并将其加载为 <a href="https://pillow.readthedocs.io/en/stable/reference/Image.html" rel="nofollow"><code>PIL.Image</code></a>，所以我们可以直接可视化：</p> <!> <div class="flex justify-center"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/butterflies_ds.png"/></div> <p>不过这些图像的尺寸各不相同，所以你需要先做预处理：</p> <ul><li><code>Resize</code> 把图像缩放到 <code>config.image_size</code> 中定义的大小。</li> <li><code>RandomHorizontalFlip</code> 通过随机水平翻转图像来做数据增强。</li> <li><code>Normalize</code> 很重要，它会把像素值缩放到 <code>[-1, 1]</code> 区间，这是模型期望的输入范围。</li></ul> <!> <p>使用 🤗 Datasets 的 <code>set_transform</code> 方法，在训练过程中按需应用 <code>preprocess</code> 函数：</p> <!> <p>你也可以再次可视化图像，确认它们已经被调整到目标尺寸。接下来，就可以把数据集封装成一个 <a href="https://pytorch.org/docs/stable/data#torch.utils.data.DataLoader" rel="nofollow">DataLoader</a> 来训练了！</p> <!> <!> <p>在 🧨 Diffusers 中，可以很方便地通过模型类和参数创建预训练模型。例如，下面创建一个 <code>UNet2DModel</code>：</p> <!> <p>通常最好先快速检查一下，样本图像的形状和模型输出形状是否一致：</p> <!> <p>很好！接下来，你还需要一个调度器为图像添加噪声。</p> <!> <p>调度器在训练和推理中的行为不同。推理时，调度器会从噪声中生成图像；训练时，调度器会取扩散过程某一步的模型输出或样本，并根据<em>噪声日程</em>与<em>更新规则</em>对图像加噪。</p> <p>我们先看看 <code>DDPMScheduler</code>，并使用 <code>add_noise</code> 方法给前面的 <code>sample_image</code> 添加一些随机噪声：</p> <!> <div class="flex justify-center"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/noisy_butterfly.png"/></div> <p>模型训练的目标，就是预测添加到图像中的噪声。当前步骤的损失可以这样计算：</p> <!> <!> <p>到这里，启动训练所需的大部分组件都准备好了，剩下的就是把它们拼起来。</p> <p>首先，你需要一个优化器和一个学习率调度器：</p> <!> <p>接着，你还需要一种评估模型的方法。评估时，我们可以使用 <code>DDPMPipeline</code> 生成一批示例图像，并把它们保存成一个网格图：</p> <!> <p>现在，你可以用 🤗 Accelerate 把这些组件包装进一个训练循环中，轻松实现 TensorBoard 日志记录、梯度累积和混合精度训练。为了把模型上传到 Hub，还需要写一个函数来创建仓库并将训练结果推送到 Hub。</p> <blockquote class="tip"><p>💡 下面的训练循环看起来可能有点长，也有点吓人，但等你真正只用一行代码启动训练时，就会觉得很值得！如果你现在只想快点开始生成图像，也可以先直接复制运行下面的代码，之后再回头仔细研究训练循环，比如等模型训练完成的时候。🤗</p></blockquote> <!> <p>呼，这段代码确实不少！不过现在你终于可以用 🤗 Accelerate 的 <code>notebook_launcher</code> 函数启动训练了。把训练循环函数、所有训练参数以及进程数（你可以改成自己可用 GPU 的数量）传进去即可：</p> <!> <p>训练完成后，来看看你的扩散模型最终生成的 🦋 蝴蝶图像 🦋 吧！</p> <!> <div class="flex justify-center"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/butterflies_final.png"/></div> <!> <p>无条件图像生成只是可训练任务中的一个例子。你可以继续访问 <a href="../training/overview">🧨 Diffusers 训练示例</a> 页面，探索更多任务和训练技术。比如：</p> <ul><li><a href="../training/text_inversion">Textual Inversion</a>：教会模型一个特定的视觉概念，并把它融入生成结果中。</li> <li><a href="../training/dreambooth">DreamBooth</a>：给定某个主体的若干输入图像，生成该主体的个性化图像。</li> <li><a href="../training/text2image">引导</a>：在你自己的数据集上微调 Stable Diffusion 模型。</li> <li><a href="../training/lora">引导</a>：使用 LoRA 这种更省内存的方法，更快地微调超大模型。</li></ul> <!> <p></p>',1);function ns(F,E){v(E,!1),O(()=>{new URLSearchParams(window.location.search).get("fw")}),X();var n=P();S("hfx23e",k=>{var Z=L();z(Z,"content",K),G(k,Z)});var M=s(D(n),2);W(M,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"});var U=s(M,2);q(U,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;",options:[{label:"Mixed",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/zh/basic_training.ipynb"},{label:"PyTorch",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/zh/pytorch/basic_training.ipynb"},{label:"TensorFlow",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/zh/tensorflow/basic_training.ipynb"},{label:"Mixed",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/zh/basic_training.ipynb"},{label:"PyTorch",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/zh/pytorch/basic_training.ipynb"},{label:"TensorFlow",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/zh/tensorflow/basic_training.ipynb"}]});var p=s(U,2);a(p,{title:"训练扩散模型",local:"训练扩散模型",headingTag:"h1"});var e=s(p,10);l(e,{code:"JTIzJTIwJUU1JUE2JTgyJUU2JTlFJTlDJUU0JUJEJUEwJUU1JTlDJUE4JTIwQ29sYWIlMjAlRTQlQjglQUQlRTglQkYlOTAlRTglQTElOEMlRUYlQkMlOEMlRTglQUYlQjclRTUlOEYlOTYlRTYlQjYlODglRTYlQjMlQTglRTklODclOEElRTYlOUQlQTUlRTUlQUUlODklRTglQTMlODUlRTYlODklODAlRTklOUMlODAlRTQlQkUlOUQlRTglQjUlOTYlMEElMjMhcGlwJTIwaW5zdGFsbCUyMGRpZmZ1c2VycyU1QnRyYWluaW5nJTVE",highlighted:`<span class="hljs-comment"># 如果你在 Colab 中运行，请取消注释来安装所需依赖</span>
<span class="hljs-comment">#!pip install diffusers[training]</span>`,lang:"py",wrap:!1});var t=s(e,4);l(t,{code:"ZnJvbSUyMGh1Z2dpbmdmYWNlX2h1YiUyMGltcG9ydCUyMG5vdGVib29rX2xvZ2luJTBBJTBBbm90ZWJvb2tfbG9naW4oKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> huggingface_hub <span class="hljs-keyword">import</span> notebook_login
<span class="hljs-meta">&gt;&gt;&gt; </span>notebook_login()`,lang:"py",wrap:!1});var c=s(t,4);l(c,{code:"aGYlMjBhdXRoJTIwbG9naW4=",highlighted:"hf auth login",lang:"bash",wrap:!1});var J=s(c,4);l(J,{code:"IXN1ZG8lMjBhcHQlMjAtcXElMjBpbnN0YWxsJTIwZ2l0LWxmcyUwQSFnaXQlMjBjb25maWclMjAtLWdsb2JhbCUyMGNyZWRlbnRpYWwuaGVscGVyJTIwc3RvcmU=",highlighted:`!<span class="hljs-built_in">sudo</span> apt -qq install git-lfs
!git config --global credential.helper store`,lang:"bash",wrap:!1});var j=s(J,2);a(j,{title:"训练配置",local:"训练配置",headingTag:"h2"});var y=s(j,4);l(y,{code:"ZnJvbSUyMGRhdGFjbGFzc2VzJTIwaW1wb3J0JTIwZGF0YWNsYXNzJTBBJTBBJTQwZGF0YWNsYXNzJTBBY2xhc3MlMjBUcmFpbmluZ0NvbmZpZyUzQSUwQSUyMCUyMCUyMCUyMGltYWdlX3NpemUlMjAlM0QlMjAxMjglMjAlMjAlMjMlMjAlRTclOTQlOUYlRTYlODglOTAlRTUlOUIlQkUlRTUlODMlOEYlRTclOUElODQlRTUlODglODYlRTglQkUlQTglRTclOEUlODclMEElMjAlMjAlMjAlMjB0cmFpbl9iYXRjaF9zaXplJTIwJTNEJTIwMTYlMEElMjAlMjAlMjAlMjBldmFsX2JhdGNoX3NpemUlMjAlM0QlMjAxNiUyMCUyMCUyMyUyMCVFOCVBRiU4NCVFNCVCQyVCMCVFNiU5NyVCNiVFNiVBRiU4RiVFNiVBQyVBMSVFOSU4NyU4NyVFNiVBMCVCNyVFNSVBNCU5QSVFNSVCMCU5MSVFNSVCQyVBMCVFNSU5QiVCRSVFNSU4MyU4RiUwQSUyMCUyMCUyMCUyMG51bV9lcG9jaHMlMjAlM0QlMjA1MCUwQSUyMCUyMCUyMCUyMGdyYWRpZW50X2FjY3VtdWxhdGlvbl9zdGVwcyUyMCUzRCUyMDElMEElMjAlMjAlMjAlMjBsZWFybmluZ19yYXRlJTIwJTNEJTIwMWUtNCUwQSUyMCUyMCUyMCUyMGxyX3dhcm11cF9zdGVwcyUyMCUzRCUyMDUwMCUwQSUyMCUyMCUyMCUyMHNhdmVfaW1hZ2VfZXBvY2hzJTIwJTNEJTIwMTAlMEElMjAlMjAlMjAlMjBzYXZlX21vZGVsX2Vwb2NocyUyMCUzRCUyMDMwJTBBJTIwJTIwJTIwJTIwbWl4ZWRfcHJlY2lzaW9uJTIwJTNEJTIwJTIyZnAxNiUyMiUyMCUyMCUyMyUyMGZsb2F0MzIlMjAlRTclOTQlQTglMjAlNjBubyU2MCVFRiVCQyU4QyVFOCU4NyVBQSVFNSU4QSVBOCVFNiVCNyVCNyVFNSU5MCU4OCVFNyVCMiVCRSVFNSVCQSVBNiVFNyU5NCVBOCUyMCU2MGZwMTYlNjAlMEElMjAlMjAlMjAlMjBvdXRwdXRfZGlyJTIwJTNEJTIwJTIyZGRwbS1idXR0ZXJmbGllcy0xMjglMjIlMjAlMjAlMjMlMjAlRTYlOUMlQUMlRTUlOUMlQjAlRTUlOTIlOEMlMjBIRiUyMEh1YiUyMCVFNCVCOCU4QSVFNyU5QSU4NCVFNiVBOCVBMSVFNSU5RSU4QiVFNSU5MCU4RCVFNyVBNyVCMCUwQSUwQSUyMCUyMCUyMCUyMHB1c2hfdG9faHViJTIwJTNEJTIwVHJ1ZSUyMCUyMCUyMyUyMCVFNiU5OCVBRiVFNSU5MCVBNiVFNSVCMCU4NiVFNCVCRiU5RCVFNSVBRCU5OCVFNSU5MCU4RSVFNyU5QSU4NCVFNiVBOCVBMSVFNSU5RSU4QiVFNCVCOCU4QSVFNCVCQyVBMCVFNSU4OCVCMCUyMEhGJTIwSHViJTBBJTIwJTIwJTIwJTIwaHViX21vZGVsX2lkJTIwJTNEJTIwJTIyJTNDeW91ci11c2VybmFtZSUzRSUyRiUzQ215LWF3ZXNvbWUtbW9kZWwlM0UlMjIlMjAlMjAlMjMlMjAlRTUlOUMlQTglMjBIRiUyMEh1YiUyMCVFNCVCOCU4QSVFNSU4OCU5QiVFNSVCQiVCQSVFNyU5QSU4NCVFNCVCQiU5MyVFNSVCQSU5MyVFNSU5MCU4RCVFNyVBNyVCMCUwQSUyMCUyMCUyMCUyMGh1Yl9wcml2YXRlX3JlcG8lMjAlM0QlMjBOb25lJTBBJTIwJTIwJTIwJTIwb3ZlcndyaXRlX291dHB1dF9kaXIlMjAlM0QlMjBUcnVlJTIwJTIwJTIzJTIwJUU5JTg3JThEJUU2JTk2JUIwJUU4JUJGJTkwJUU4JUExJThDJTIwbm90ZWJvb2slMjAlRTYlOTclQjYlRTYlOTglQUYlRTUlOTAlQTYlRTglQTYlODYlRTclOUIlOTYlRTYlOTclQTclRTYlQTglQTElRTUlOUUlOEIlMEElMjAlMjAlMjAlMjBzZWVkJTIwJTNEJTIwMCUwQSUwQSUwQWNvbmZpZyUyMCUzRCUyMFRyYWluaW5nQ29uZmlnKCk=",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> dataclasses <span class="hljs-keyword">import</span> dataclass
<span class="hljs-meta">&gt;&gt;&gt; </span>@dataclass
<span class="hljs-meta">... </span><span class="hljs-keyword">class</span> <span class="hljs-title class_">TrainingConfig</span>:
<span class="hljs-meta">... </span> image_size = <span class="hljs-number">128</span> <span class="hljs-comment"># 生成图像的分辨率</span>
<span class="hljs-meta">... </span> train_batch_size = <span class="hljs-number">16</span>
<span class="hljs-meta">... </span> eval_batch_size = <span class="hljs-number">16</span> <span class="hljs-comment"># 评估时每次采样多少张图像</span>
<span class="hljs-meta">... </span> num_epochs = <span class="hljs-number">50</span>
<span class="hljs-meta">... </span> gradient_accumulation_steps = <span class="hljs-number">1</span>
<span class="hljs-meta">... </span> learning_rate = <span class="hljs-number">1e-4</span>
<span class="hljs-meta">... </span> lr_warmup_steps = <span class="hljs-number">500</span>
<span class="hljs-meta">... </span> save_image_epochs = <span class="hljs-number">10</span>
<span class="hljs-meta">... </span> save_model_epochs = <span class="hljs-number">30</span>
<span class="hljs-meta">... </span> mixed_precision = <span class="hljs-string">&quot;fp16&quot;</span> <span class="hljs-comment"># float32 用 \`no\`，自动混合精度用 \`fp16\`</span>
<span class="hljs-meta">... </span> output_dir = <span class="hljs-string">&quot;ddpm-butterflies-128&quot;</span> <span class="hljs-comment"># 本地和 HF Hub 上的模型名称</span>
<span class="hljs-meta">... </span> push_to_hub = <span class="hljs-literal">True</span> <span class="hljs-comment"># 是否将保存后的模型上传到 HF Hub</span>
<span class="hljs-meta">... </span> hub_model_id = <span class="hljs-string">&quot;&lt;your-username&gt;/&lt;my-awesome-model&gt;&quot;</span> <span class="hljs-comment"># 在 HF Hub 上创建的仓库名称</span>
<span class="hljs-meta">... </span> hub_private_repo = <span class="hljs-literal">None</span>
<span class="hljs-meta">... </span> overwrite_output_dir = <span class="hljs-literal">True</span> <span class="hljs-comment"># 重新运行 notebook 时是否覆盖旧模型</span>
<span class="hljs-meta">... </span> seed = <span class="hljs-number">0</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>config = TrainingConfig()`,lang:"py",wrap:!1});var o=s(y,2);a(o,{title:"加载数据集",local:"加载数据集",headingTag:"h2"});var T=s(o,4);l(T,{code:"ZnJvbSUyMGRhdGFzZXRzJTIwaW1wb3J0JTIwbG9hZF9kYXRhc2V0JTBBJTBBY29uZmlnLmRhdGFzZXRfbmFtZSUyMCUzRCUyMCUyMmh1Z2dhbiUyRnNtaXRoc29uaWFuX2J1dHRlcmZsaWVzX3N1YnNldCUyMiUwQWRhdGFzZXQlMjAlM0QlMjBsb2FkX2RhdGFzZXQoY29uZmlnLmRhdGFzZXRfbmFtZSUyQyUyMHNwbGl0JTNEJTIydHJhaW4lMjIp",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> load_dataset
<span class="hljs-meta">&gt;&gt;&gt; </span>config.dataset_name = <span class="hljs-string">&quot;huggan/smithsonian_butterflies_subset&quot;</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>dataset = load_dataset(config.dataset_name, split=<span class="hljs-string">&quot;train&quot;</span>)`,lang:"py",wrap:!1});var h=s(T,6);l(h,{code:"aW1wb3J0JTIwbWF0cGxvdGxpYi5weXBsb3QlMjBhcyUyMHBsdCUwQSUwQWZpZyUyQyUyMGF4cyUyMCUzRCUyMHBsdC5zdWJwbG90cygxJTJDJTIwNCUyQyUyMGZpZ3NpemUlM0QoMTYlMkMlMjA0KSklMEFmb3IlMjBpJTJDJTIwaW1hZ2UlMjBpbiUyMGVudW1lcmF0ZShkYXRhc2V0JTVCJTNBNCU1RCU1QiUyMmltYWdlJTIyJTVEKSUzQSUwQSUyMCUyMCUyMCUyMGF4cyU1QmklNUQuaW1zaG93KGltYWdlKSUwQSUyMCUyMCUyMCUyMGF4cyU1QmklNUQuc2V0X2F4aXNfb2ZmKCklMEFmaWcuc2hvdygp",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> matplotlib.pyplot <span class="hljs-keyword">as</span> plt
<span class="hljs-meta">&gt;&gt;&gt; </span>fig, axs = plt.subplots(<span class="hljs-number">1</span>, <span class="hljs-number">4</span>, figsize=(<span class="hljs-number">16</span>, <span class="hljs-number">4</span>))
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">for</span> i, image <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(dataset[:<span class="hljs-number">4</span>][<span class="hljs-string">&quot;image&quot;</span>]):
<span class="hljs-meta">... </span> axs[i].imshow(image)
<span class="hljs-meta">... </span> axs[i].set_axis_off()
<span class="hljs-meta">&gt;&gt;&gt; </span>fig.show()`,lang:"py",wrap:!1});var i=s(h,8);l(i,{code:"ZnJvbSUyMHRvcmNodmlzaW9uJTIwaW1wb3J0JTIwdHJhbnNmb3JtcyUwQSUwQXByZXByb2Nlc3MlMjAlM0QlMjB0cmFuc2Zvcm1zLkNvbXBvc2UoJTBBJTIwJTIwJTIwJTIwJTVCJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwdHJhbnNmb3Jtcy5SZXNpemUoKGNvbmZpZy5pbWFnZV9zaXplJTJDJTIwY29uZmlnLmltYWdlX3NpemUpKSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHRyYW5zZm9ybXMuUmFuZG9tSG9yaXpvbnRhbEZsaXAoKSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHRyYW5zZm9ybXMuVG9UZW5zb3IoKSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHRyYW5zZm9ybXMuTm9ybWFsaXplKCU1QjAuNSU1RCUyQyUyMCU1QjAuNSU1RCklMkMlMEElMjAlMjAlMjAlMjAlNUQlMEEp",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> torchvision <span class="hljs-keyword">import</span> transforms
<span class="hljs-meta">&gt;&gt;&gt; </span>preprocess = transforms.Compose(
<span class="hljs-meta">... </span> [
<span class="hljs-meta">... </span> transforms.Resize((config.image_size, config.image_size)),
<span class="hljs-meta">... </span> transforms.RandomHorizontalFlip(),
<span class="hljs-meta">... </span> transforms.ToTensor(),
<span class="hljs-meta">... </span> transforms.Normalize([<span class="hljs-number">0.5</span>], [<span class="hljs-number">0.5</span>]),
<span class="hljs-meta">... </span> ]
<span class="hljs-meta">... </span>)`,lang:"py",wrap:!1});var m=s(i,4);l(m,{code:"ZGVmJTIwdHJhbnNmb3JtKGV4YW1wbGVzKSUzQSUwQSUyMCUyMCUyMCUyMGltYWdlcyUyMCUzRCUyMCU1QnByZXByb2Nlc3MoaW1hZ2UuY29udmVydCglMjJSR0IlMjIpKSUyMGZvciUyMGltYWdlJTIwaW4lMjBleGFtcGxlcyU1QiUyMmltYWdlJTIyJTVEJTVEJTBBJTIwJTIwJTIwJTIwcmV0dXJuJTIwJTdCJTIyaW1hZ2VzJTIyJTNBJTIwaW1hZ2VzJTdEJTBBJTBBJTBBZGF0YXNldC5zZXRfdHJhbnNmb3JtKHRyYW5zZm9ybSk=",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">def</span> <span class="hljs-title function_">transform</span>(<span class="hljs-params">examples</span>):
<span class="hljs-meta">... </span> images = [preprocess(image.convert(<span class="hljs-string">&quot;RGB&quot;</span>)) <span class="hljs-keyword">for</span> image <span class="hljs-keyword">in</span> examples[<span class="hljs-string">&quot;image&quot;</span>]]
<span class="hljs-meta">... </span> <span class="hljs-keyword">return</span> {<span class="hljs-string">&quot;images&quot;</span>: images}
<span class="hljs-meta">&gt;&gt;&gt; </span>dataset.set_transform(transform)`,lang:"py",wrap:!1});var w=s(m,4);l(w,{code:"aW1wb3J0JTIwdG9yY2glMEElMEF0cmFpbl9kYXRhbG9hZGVyJTIwJTNEJTIwdG9yY2gudXRpbHMuZGF0YS5EYXRhTG9hZGVyKGRhdGFzZXQlMkMlMjBiYXRjaF9zaXplJTNEY29uZmlnLnRyYWluX2JhdGNoX3NpemUlMkMlMjBzaHVmZmxlJTNEVHJ1ZSk=",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> torch
<span class="hljs-meta">&gt;&gt;&gt; </span>train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=<span class="hljs-literal">True</span>)`,lang:"py",wrap:!1});var C=s(w,2);a(C,{title:"创建 UNet2DModel",local:"创建-unet2dmodel",headingTag:"h2"});var r=s(C,4);l(r,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMFVOZXQyRE1vZGVsJTBBJTBBbW9kZWwlMjAlM0QlMjBVTmV0MkRNb2RlbCglMEElMjAlMjAlMjAlMjBzYW1wbGVfc2l6ZSUzRGNvbmZpZy5pbWFnZV9zaXplJTJDJTIwJTIwJTIzJTIwJUU3JTlCJUFFJUU2JUEwJTg3JUU1JTlCJUJFJUU1JTgzJThGJUU1JTg4JTg2JUU4JUJFJUE4JUU3JThFJTg3JTBBJTIwJTIwJTIwJTIwaW5fY2hhbm5lbHMlM0QzJTJDJTIwJTIwJTIzJTIwJUU4JUJFJTkzJUU1JTg1JUE1JUU5JTgwJTlBJUU5JTgxJTkzJUU2JTk1JUIwJUVGJUJDJThDUkdCJTIwJUU1JTlCJUJFJUU1JTgzJThGJUU0JUI4JUJBJTIwMyUwQSUyMCUyMCUyMCUyMG91dF9jaGFubmVscyUzRDMlMkMlMjAlMjAlMjMlMjAlRTglQkUlOTMlRTUlODclQkElRTklODAlOUElRTklODElOTMlRTYlOTUlQjAlMEElMjAlMjAlMjAlMjBsYXllcnNfcGVyX2Jsb2NrJTNEMiUyQyUyMCUyMCUyMyUyMCVFNiVBRiU4RiVFNCVCOCVBQSUyMFVOZXQlMjBibG9jayUyMCVFNCVCOCVBRCVFNCVCRCVCRiVFNyU5NCVBOCVFNSVBNCU5QSVFNSVCMCU5MSVFNCVCOCVBQSUyMFJlc05ldCUyMCVFNSVCMSU4MiUwQSUyMCUyMCUyMCUyMGJsb2NrX291dF9jaGFubmVscyUzRCgxMjglMkMlMjAxMjglMkMlMjAyNTYlMkMlMjAyNTYlMkMlMjA1MTIlMkMlMjA1MTIpJTJDJTIwJTIwJTIzJTIwJUU2JUFGJThGJUU0JUI4JUFBJTIwVU5ldCUyMGJsb2NrJTIwJUU3JTlBJTg0JUU4JUJFJTkzJUU1JTg3JUJBJUU5JTgwJTlBJUU5JTgxJTkzJUU2JTk1JUIwJTBBJTIwJTIwJTIwJTIwZG93bl9ibG9ja190eXBlcyUzRCglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJEb3duQmxvY2syRCUyMiUyQyUyMCUyMCUyMyUyMCVFNiVBMCU4NyVFNSU4NyU4NiVFNyU5QSU4NCUyMFJlc05ldCUyMCVFNCVCOCU4QiVFOSU4NyU4NyVFNiVBMCVCNyVFNSU5RCU5NyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMkRvd25CbG9jazJEJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIyRG93bkJsb2NrMkQlMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJEb3duQmxvY2syRCUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMkF0dG5Eb3duQmxvY2syRCUyMiUyQyUyMCUyMCUyMyUyMCVFNSVCOCVBNiVFNyVBOSVCQSVFOSU5NyVCNCVFOCU4NyVBQSVFNiVCMyVBOCVFNiU4NCU4RiVFNSU4QSU5QiVFNyU5QSU4NCUyMFJlc05ldCUyMCVFNCVCOCU4QiVFOSU4NyU4NyVFNiVBMCVCNyVFNSU5RCU5NyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMkRvd25CbG9jazJEJTIyJTJDJTBBJTIwJTIwJTIwJTIwKSUyQyUwQSUyMCUyMCUyMCUyMHVwX2Jsb2NrX3R5cGVzJTNEKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMlVwQmxvY2syRCUyMiUyQyUyMCUyMCUyMyUyMCVFNiVBMCU4NyVFNSU4NyU4NiVFNyU5QSU4NCUyMFJlc05ldCUyMCVFNCVCOCU4QSVFOSU4NyU4NyVFNiVBMCVCNyVFNSU5RCU5NyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMkF0dG5VcEJsb2NrMkQlMjIlMkMlMjAlMjAlMjMlMjAlRTUlQjglQTYlRTclQTklQkElRTklOTclQjQlRTglODclQUElRTYlQjMlQTglRTYlODQlOEYlRTUlOEElOUIlRTclOUElODQlMjBSZXNOZXQlMjAlRTQlQjglOEElRTklODclODclRTYlQTAlQjclRTUlOUQlOTclMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJVcEJsb2NrMkQlMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJVcEJsb2NrMkQlMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJVcEJsb2NrMkQlMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJVcEJsb2NrMkQlMjIlMkMlMEElMjAlMjAlMjAlMjApJTJDJTBBKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> UNet2DModel
<span class="hljs-meta">&gt;&gt;&gt; </span>model = UNet2DModel(
<span class="hljs-meta">... </span> sample_size=config.image_size, <span class="hljs-comment"># 目标图像分辨率</span>
<span class="hljs-meta">... </span> in_channels=<span class="hljs-number">3</span>, <span class="hljs-comment"># 输入通道数，RGB 图像为 3</span>
<span class="hljs-meta">... </span> out_channels=<span class="hljs-number">3</span>, <span class="hljs-comment"># 输出通道数</span>
<span class="hljs-meta">... </span> layers_per_block=<span class="hljs-number">2</span>, <span class="hljs-comment"># 每个 UNet block 中使用多少个 ResNet 层</span>
<span class="hljs-meta">... </span> block_out_channels=(<span class="hljs-number">128</span>, <span class="hljs-number">128</span>, <span class="hljs-number">256</span>, <span class="hljs-number">256</span>, <span class="hljs-number">512</span>, <span class="hljs-number">512</span>), <span class="hljs-comment"># 每个 UNet block 的输出通道数</span>
<span class="hljs-meta">... </span> down_block_types=(
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;DownBlock2D&quot;</span>, <span class="hljs-comment"># 标准的 ResNet 下采样块</span>
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;DownBlock2D&quot;</span>,
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;DownBlock2D&quot;</span>,
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;DownBlock2D&quot;</span>,
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;AttnDownBlock2D&quot;</span>, <span class="hljs-comment"># 带空间自注意力的 ResNet 下采样块</span>
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;DownBlock2D&quot;</span>,
<span class="hljs-meta">... </span> ),
<span class="hljs-meta">... </span> up_block_types=(
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;UpBlock2D&quot;</span>, <span class="hljs-comment"># 标准的 ResNet 上采样块</span>
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;AttnUpBlock2D&quot;</span>, <span class="hljs-comment"># 带空间自注意力的 ResNet 上采样块</span>
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;UpBlock2D&quot;</span>,
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;UpBlock2D&quot;</span>,
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;UpBlock2D&quot;</span>,
<span class="hljs-meta">... </span> <span class="hljs-string">&quot;UpBlock2D&quot;</span>,
<span class="hljs-meta">... </span> ),
<span class="hljs-meta">... </span>)`,lang:"py",wrap:!1});var g=s(r,4);l(g,{code:"c2FtcGxlX2ltYWdlJTIwJTNEJTIwZGF0YXNldCU1QjAlNUQlNUIlMjJpbWFnZXMlMjIlNUQudW5zcXVlZXplKDApJTBBcHJpbnQoJTIySW5wdXQlMjBzaGFwZSUzQSUyMiUyQyUyMHNhbXBsZV9pbWFnZS5zaGFwZSklMEElMEFwcmludCglMjJPdXRwdXQlMjBzaGFwZSUzQSUyMiUyQyUyMG1vZGVsKHNhbXBsZV9pbWFnZSUyQyUyMHRpbWVzdGVwJTNEMCkuc2FtcGxlLnNoYXBlKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span>sample_image = dataset[<span class="hljs-number">0</span>][<span class="hljs-string">&quot;images&quot;</span>].unsqueeze(<span class="hljs-number">0</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;Input shape:&quot;</span>, sample_image.shape)
Input shape: torch.Size([<span class="hljs-number">1</span>, <span class="hljs-number">3</span>, <span class="hljs-number">128</span>, <span class="hljs-number">128</span>])
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;Output shape:&quot;</span>, model(sample_image, timestep=<span class="hljs-number">0</span>).sample.shape)
Output shape: torch.Size([<span class="hljs-number">1</span>, <span class="hljs-number">3</span>, <span class="hljs-number">128</span>, <span class="hljs-number">128</span>])`,lang:"py",wrap:!1});var d=s(g,4);a(d,{title:"创建调度器",local:"创建调度器",headingTag:"h2"});var b=s(d,6);l(b,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwUElMJTIwaW1wb3J0JTIwSW1hZ2UlMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwRERQTVNjaGVkdWxlciUwQSUwQW5vaXNlX3NjaGVkdWxlciUyMCUzRCUyMEREUE1TY2hlZHVsZXIobnVtX3RyYWluX3RpbWVzdGVwcyUzRDEwMDApJTBBbm9pc2UlMjAlM0QlMjB0b3JjaC5yYW5kbihzYW1wbGVfaW1hZ2Uuc2hhcGUpJTBBdGltZXN0ZXBzJTIwJTNEJTIwdG9yY2guTG9uZ1RlbnNvciglNUI1MCU1RCklMEFub2lzeV9pbWFnZSUyMCUzRCUyMG5vaXNlX3NjaGVkdWxlci5hZGRfbm9pc2Uoc2FtcGxlX2ltYWdlJTJDJTIwbm9pc2UlMkMlMjB0aW1lc3RlcHMpJTBBJTBBSW1hZ2UuZnJvbWFycmF5KCgobm9pc3lfaW1hZ2UucGVybXV0ZSgwJTJDJTIwMiUyQyUyMDMlMkMlMjAxKSUyMCUyQiUyMDEuMCklMjAqJTIwMTI3LjUpLnR5cGUodG9yY2gudWludDgpLm51bXB5KCklNUIwJTVEKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> torch
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> PIL <span class="hljs-keyword">import</span> Image
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DDPMScheduler
<span class="hljs-meta">&gt;&gt;&gt; </span>noise_scheduler = DDPMScheduler(num_train_timesteps=<span class="hljs-number">1000</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>noise = torch.randn(sample_image.shape)
<span class="hljs-meta">&gt;&gt;&gt; </span>timesteps = torch.LongTensor([<span class="hljs-number">50</span>])
<span class="hljs-meta">&gt;&gt;&gt; </span>noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)
<span class="hljs-meta">&gt;&gt;&gt; </span>Image.fromarray(((noisy_image.permute(<span class="hljs-number">0</span>, <span class="hljs-number">2</span>, <span class="hljs-number">3</span>, <span class="hljs-number">1</span>) + <span class="hljs-number">1.0</span>) * <span class="hljs-number">127.5</span>).<span class="hljs-built_in">type</span>(torch.uint8).numpy()[<span class="hljs-number">0</span>])`,lang:"py",wrap:!1});var I=s(b,6);l(I,{code:"aW1wb3J0JTIwdG9yY2gubm4uZnVuY3Rpb25hbCUyMGFzJTIwRiUwQSUwQW5vaXNlX3ByZWQlMjAlM0QlMjBtb2RlbChub2lzeV9pbWFnZSUyQyUyMHRpbWVzdGVwcykuc2FtcGxlJTBBbG9zcyUyMCUzRCUyMEYubXNlX2xvc3Mobm9pc2VfcHJlZCUyQyUyMG5vaXNlKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> torch.nn.functional <span class="hljs-keyword">as</span> F
<span class="hljs-meta">&gt;&gt;&gt; </span>noise_pred = model(noisy_image, timesteps).sample
<span class="hljs-meta">&gt;&gt;&gt; </span>loss = F.mse_loss(noise_pred, noise)`,lang:"py",wrap:!1});var u=s(I,2);a(u,{title:"训练模型",local:"训练模型",headingTag:"h2"});var A=s(u,6);l(A,{code:"ZnJvbSUyMGRpZmZ1c2Vycy5vcHRpbWl6YXRpb24lMjBpbXBvcnQlMjBnZXRfY29zaW5lX3NjaGVkdWxlX3dpdGhfd2FybXVwJTBBJTBBb3B0aW1pemVyJTIwJTNEJTIwdG9yY2gub3B0aW0uQWRhbVcobW9kZWwucGFyYW1ldGVycygpJTJDJTIwbHIlM0Rjb25maWcubGVhcm5pbmdfcmF0ZSklMEFscl9zY2hlZHVsZXIlMjAlM0QlMjBnZXRfY29zaW5lX3NjaGVkdWxlX3dpdGhfd2FybXVwKCUwQSUyMCUyMCUyMCUyMG9wdGltaXplciUzRG9wdGltaXplciUyQyUwQSUyMCUyMCUyMCUyMG51bV93YXJtdXBfc3RlcHMlM0Rjb25maWcubHJfd2FybXVwX3N0ZXBzJTJDJTBBJTIwJTIwJTIwJTIwbnVtX3RyYWluaW5nX3N0ZXBzJTNEKGxlbih0cmFpbl9kYXRhbG9hZGVyKSUyMColMjBjb25maWcubnVtX2Vwb2NocyklMkMlMEEp",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> diffusers.optimization <span class="hljs-keyword">import</span> get_cosine_schedule_with_warmup
<span class="hljs-meta">&gt;&gt;&gt; </span>optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
<span class="hljs-meta">&gt;&gt;&gt; </span>lr_scheduler = get_cosine_schedule_with_warmup(
<span class="hljs-meta">... </span> optimizer=optimizer,
<span class="hljs-meta">... </span> num_warmup_steps=config.lr_warmup_steps,
<span class="hljs-meta">... </span> num_training_steps=(<span class="hljs-built_in">len</span>(train_dataloader) * config.num_epochs),
<span class="hljs-meta">... </span>)`,lang:"py",wrap:!1});var V=s(A,4);l(V,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMEREUE1QaXBlbGluZSUwQWZyb20lMjBkaWZmdXNlcnMudXRpbHMlMjBpbXBvcnQlMjBtYWtlX2ltYWdlX2dyaWQlMEFpbXBvcnQlMjBvcyUwQSUwQWRlZiUyMGV2YWx1YXRlKGNvbmZpZyUyQyUyMGVwb2NoJTJDJTIwcGlwZWxpbmUpJTNBJTBBJTIwJTIwJTIwJTIwJTIzJTIwJUU0JUJCJThFJUU5JTlBJThGJUU2JTlDJUJBJUU1JTk5JUFBJUU1JUEzJUIwJUU5JTg3JTg3JUU2JUEwJUI3JUU1JTlCJUJFJUU1JTgzJThGJUVGJUJDJTg4JUU4JUJGJTk5JUU1JUIwJUIxJUU2JTk4JUFGJUU1JThGJThEJUU1JTkwJTkxJUU2JTg5JUE5JUU2JTk1JUEzJUU4JUJGJTg3JUU3JUE4JThCJUVGJUJDJTg5JTBBJTIwJTIwJTIwJTIwJTIzJTIwJUU3JUFFJUExJUU5JTgxJTkzJUU5JUJCJTk4JUU4JUFFJUE0JUU4JUJFJTkzJUU1JTg3JUJBJUU3JUIxJUJCJUU1JTlFJThCJUU2JTk4JUFGJTIwJTYwTGlzdCU1QlBJTC5JbWFnZSU1RCU2MCUwQSUyMCUyMCUyMCUyMGltYWdlcyUyMCUzRCUyMHBpcGVsaW5lKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJhdGNoX3NpemUlM0Rjb25maWcuZXZhbF9iYXRjaF9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZ2VuZXJhdG9yJTNEdG9yY2guR2VuZXJhdG9yKGRldmljZSUzRCdjcHUnKS5tYW51YWxfc2VlZChjb25maWcuc2VlZCklMkMlMjAlMjMlMjAlRTUlOEQlOTUlRTclOEIlQUMlRTQlQkQlQkYlRTclOTQlQTglRTQlQjglODAlRTQlQjglQUElMjB0b3JjaCUyMGdlbmVyYXRvciVFRiVCQyU4QyVFOSU4MSVCRiVFNSU4NSU4RCVFNSU5QiU5RSVFOSU4MCU4MCVFNCVCOCVCQiVFOCVBRSVBRCVFNyVCQiU4MyVFNSVCRSVBQSVFNyU4RSVBRiVFNyU5QSU4NCVFOSU5QSU4RiVFNiU5QyVCQSVFNyU4QSVCNiVFNiU4MCU4MSUwQSUyMCUyMCUyMCUyMCkuaW1hZ2VzJTBBJTBBJTIwJTIwJTIwJTIwJTIzJTIwJUU2JThBJThBJUU1JTlCJUJFJUU1JTgzJThGJUU2JThCJUJDJUU2JTg4JTkwJUU3JUJEJTkxJUU2JUEwJUJDJTBBJTIwJTIwJTIwJTIwaW1hZ2VfZ3JpZCUyMCUzRCUyMG1ha2VfaW1hZ2VfZ3JpZChpbWFnZXMlMkMlMjByb3dzJTNENCUyQyUyMGNvbHMlM0Q0KSUwQSUwQSUyMCUyMCUyMCUyMCUyMyUyMCVFNCVCRiU5RCVFNSVBRCU5OCVFNSU5QiVCRSVFNSU4MyU4RiUwQSUyMCUyMCUyMCUyMHRlc3RfZGlyJTIwJTNEJTIwb3MucGF0aC5qb2luKGNvbmZpZy5vdXRwdXRfZGlyJTJDJTIwJTIyc2FtcGxlcyUyMiklMEElMjAlMjAlMjAlMjBvcy5tYWtlZGlycyh0ZXN0X2RpciUyQyUyMGV4aXN0X29rJTNEVHJ1ZSklMEElMjAlMjAlMjAlMjBpbWFnZV9ncmlkLnNhdmUoZiUyMiU3QnRlc3RfZGlyJTdEJTJGJTdCZXBvY2glM0EwNGQlN0QucG5nJTIyKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DDPMPipeline
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> diffusers.utils <span class="hljs-keyword">import</span> make_image_grid
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> os
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">def</span> <span class="hljs-title function_">evaluate</span>(<span class="hljs-params">config, epoch, pipeline</span>):
<span class="hljs-meta">... </span> <span class="hljs-comment"># 从随机噪声采样图像（这就是反向扩散过程）</span>
<span class="hljs-meta">... </span> <span class="hljs-comment"># 管道默认输出类型是 \`List[PIL.Image]\`</span>
<span class="hljs-meta">... </span> images = pipeline(
<span class="hljs-meta">... </span> batch_size=config.eval_batch_size,
<span class="hljs-meta">... </span> generator=torch.Generator(device=<span class="hljs-string">&#x27;cpu&#x27;</span>).manual_seed(config.seed), <span class="hljs-comment"># 单独使用一个 torch generator，避免回退主训练循环的随机状态</span>
<span class="hljs-meta">... </span> ).images
<span class="hljs-meta">... </span> <span class="hljs-comment"># 把图像拼成网格</span>
<span class="hljs-meta">... </span> image_grid = make_image_grid(images, rows=<span class="hljs-number">4</span>, cols=<span class="hljs-number">4</span>)
<span class="hljs-meta">... </span> <span class="hljs-comment"># 保存图像</span>
<span class="hljs-meta">... </span> test_dir = os.path.join(config.output_dir, <span class="hljs-string">&quot;samples&quot;</span>)
<span class="hljs-meta">... </span> os.makedirs(test_dir, exist_ok=<span class="hljs-literal">True</span>)
<span class="hljs-meta">... </span> image_grid.save(<span class="hljs-string">f&quot;<span class="hljs-subst">{test_dir}</span>/<span class="hljs-subst">{epoch:04d}</span>.png&quot;</span>)`,lang:"py",wrap:!1});var Q=s(V,6);l(Q,{code:"ZnJvbSUyMGFjY2VsZXJhdGUlMjBpbXBvcnQlMjBBY2NlbGVyYXRvciUwQWZyb20lMjBodWdnaW5nZmFjZV9odWIlMjBpbXBvcnQlMjBjcmVhdGVfcmVwbyUyQyUyMHVwbG9hZF9mb2xkZXIlMEFmcm9tJTIwdHFkbS5hdXRvJTIwaW1wb3J0JTIwdHFkbSUwQWZyb20lMjBwYXRobGliJTIwaW1wb3J0JTIwUGF0aCUwQWltcG9ydCUyMG9zJTBBJTBBZGVmJTIwdHJhaW5fbG9vcChjb25maWclMkMlMjBtb2RlbCUyQyUyMG5vaXNlX3NjaGVkdWxlciUyQyUyMG9wdGltaXplciUyQyUyMHRyYWluX2RhdGFsb2FkZXIlMkMlMjBscl9zY2hlZHVsZXIpJTNBJTBBJTIwJTIwJTIwJTIwJTIzJTIwJUU1JTg4JTlEJUU1JUE3JThCJUU1JThDJTk2JTIwYWNjZWxlcmF0b3IlMjAlRTUlOTIlOEMlMjB0ZW5zb3Jib2FyZCUyMCVFNiU5NyVBNSVFNSVCRiU5NyUwQSUyMCUyMCUyMCUyMGFjY2VsZXJhdG9yJTIwJTNEJTIwQWNjZWxlcmF0b3IoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwbWl4ZWRfcHJlY2lzaW9uJTNEY29uZmlnLm1peGVkX3ByZWNpc2lvbiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdyYWRpZW50X2FjY3VtdWxhdGlvbl9zdGVwcyUzRGNvbmZpZy5ncmFkaWVudF9hY2N1bXVsYXRpb25fc3RlcHMlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBsb2dfd2l0aCUzRCUyMnRlbnNvcmJvYXJkJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwcHJvamVjdF9kaXIlM0Rvcy5wYXRoLmpvaW4oY29uZmlnLm91dHB1dF9kaXIlMkMlMjAlMjJsb2dzJTIyKSUyQyUwQSUyMCUyMCUyMCUyMCklMEElMjAlMjAlMjAlMjBpZiUyMGFjY2VsZXJhdG9yLmlzX21haW5fcHJvY2VzcyUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlmJTIwY29uZmlnLm91dHB1dF9kaXIlMjBpcyUyMG5vdCUyME5vbmUlM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBvcy5tYWtlZGlycyhjb25maWcub3V0cHV0X2RpciUyQyUyMGV4aXN0X29rJTNEVHJ1ZSklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBpZiUyMGNvbmZpZy5wdXNoX3RvX2h1YiUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHJlcG9faWQlMjAlM0QlMjBjcmVhdGVfcmVwbyglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjByZXBvX2lkJTNEY29uZmlnLmh1Yl9tb2RlbF9pZCUyMG9yJTIwUGF0aChjb25maWcub3V0cHV0X2RpcikubmFtZSUyQyUyMGV4aXN0X29rJTNEVHJ1ZSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCkucmVwb19pZCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGFjY2VsZXJhdG9yLmluaXRfdHJhY2tlcnMoJTIydHJhaW5fZXhhbXBsZSUyMiklMEElMEElMjAlMjAlMjAlMjAlMjMlMjAlRTUlODclODYlRTUlQTQlODclRTYlODklODAlRTYlOUMlODklRTUlQUYlQjklRTglQjElQTElMEElMjAlMjAlMjAlMjAlMjMlMjAlRTQlQjglOEQlRTklOUMlODAlRTglQTYlODElRTglQUUlQjAlRTQlQkQlOEYlRTUlOUIlQkElRTUlQUUlOUElRTklQTElQkElRTUlQkElOEYlRUYlQkMlOEMlRTUlOEYlQUElRTglQTYlODElRTglQTclQTMlRTUlOEMlODUlRTYlOTclQjYlRTUlOTIlOEMlRTQlQkMlQTAlRTclQkIlOTklMjBwcmVwYXJlJTIwJUU3JTlBJTg0JUU5JUExJUJBJUU1JUJBJThGJUU0JUI4JTgwJUU4JTg3JUI0JUU1JThEJUIzJUU1JThGJUFGJUUzJTgwJTgyJTBBJTIwJTIwJTIwJTIwbW9kZWwlMkMlMjBvcHRpbWl6ZXIlMkMlMjB0cmFpbl9kYXRhbG9hZGVyJTJDJTIwbHJfc2NoZWR1bGVyJTIwJTNEJTIwYWNjZWxlcmF0b3IucHJlcGFyZSglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBtb2RlbCUyQyUyMG9wdGltaXplciUyQyUyMHRyYWluX2RhdGFsb2FkZXIlMkMlMjBscl9zY2hlZHVsZXIlMEElMjAlMjAlMjAlMjApJTBBJTBBJTIwJTIwJTIwJTIwZ2xvYmFsX3N0ZXAlMjAlM0QlMjAwJTBBJTBBJTIwJTIwJTIwJTIwJTIzJTIwJUU1JUJDJTgwJUU1JUE3JThCJUU4JUFFJUFEJUU3JUJCJTgzJUU2JUE4JUExJUU1JTlFJThCJTBBJTIwJTIwJTIwJTIwZm9yJTIwZXBvY2glMjBpbiUyMHJhbmdlKGNvbmZpZy5udW1fZXBvY2hzKSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHByb2dyZXNzX2JhciUyMCUzRCUyMHRxZG0odG90YWwlM0RsZW4odHJhaW5fZGF0YWxvYWRlciklMkMlMjBkaXNhYmxlJTNEbm90JTIwYWNjZWxlcmF0b3IuaXNfbG9jYWxfbWFpbl9wcm9jZXNzKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHByb2dyZXNzX2Jhci5zZXRfZGVzY3JpcHRpb24oZiUyMkVwb2NoJTIwJTdCZXBvY2glN0QlMjIpJTBBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZm9yJTIwc3RlcCUyQyUyMGJhdGNoJTIwaW4lMjBlbnVtZXJhdGUodHJhaW5fZGF0YWxvYWRlciklM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBjbGVhbl9pbWFnZXMlMjAlM0QlMjBiYXRjaCU1QiUyMmltYWdlcyUyMiU1RCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMCVFNCVCOCVCQSVFNSU5QiVCRSVFNSU4MyU4RiVFOSU4NyU4NyVFNiVBMCVCNyVFNSU5OSVBQSVFNSVBMyVCMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG5vaXNlJTIwJTNEJTIwdG9yY2gucmFuZG4oY2xlYW5faW1hZ2VzLnNoYXBlJTJDJTIwZGV2aWNlJTNEY2xlYW5faW1hZ2VzLmRldmljZSklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBicyUyMCUzRCUyMGNsZWFuX2ltYWdlcy5zaGFwZSU1QjAlNUQlMEElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjAlRTQlQjglQkElRTYlQUYlOEYlRTUlQkMlQTAlRTUlOUIlQkUlRTUlODMlOEYlRTklOUElOEYlRTYlOUMlQkElRTklODclODclRTYlQTAlQjclRTQlQjglODAlRTQlQjglQUElRTYlOTclQjYlRTklOTclQjQlRTYlQUQlQTUlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjB0aW1lc3RlcHMlMjAlM0QlMjB0b3JjaC5yYW5kaW50KCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMDAlMkMlMjBub2lzZV9zY2hlZHVsZXIuY29uZmlnLm51bV90cmFpbl90aW1lc3RlcHMlMkMlMjAoYnMlMkMpJTJDJTIwZGV2aWNlJTNEY2xlYW5faW1hZ2VzLmRldmljZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGR0eXBlJTNEdG9yY2guaW50NjQlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjApJTBBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwJUU2JThDJTg5JUU3JTg1JUE3JUU2JUFGJThGJUU0JUI4JUFBJUU2JTk3JUI2JUU5JTk3JUI0JUU2JUFEJUE1JUU1JUFGJUI5JUU1JUJBJTk0JUU3JTlBJTg0JUU1JTk5JUFBJUU1JUEzJUIwJUU1JUJDJUJBJUU1JUJBJUE2JUU3JUJCJTk5JUU1JUI5JUIyJUU1JTg3JTgwJUU1JTlCJUJFJUU1JTgzJThGJUU1JThBJUEwJUU1JTk5JUFBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwJUVGJUJDJTg4JUU4JUJGJTk5JUU1JUIwJUIxJUU2JTk4JUFGJUU1JTg5JThEJUU1JTkwJTkxJUU2JTg5JUE5JUU2JTk1JUEzJUU4JUJGJTg3JUU3JUE4JThCJUVGJUJDJTg5JTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwbm9pc3lfaW1hZ2VzJTIwJTNEJTIwbm9pc2Vfc2NoZWR1bGVyLmFkZF9ub2lzZShjbGVhbl9pbWFnZXMlMkMlMjBub2lzZSUyQyUyMHRpbWVzdGVwcyklMEElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjB3aXRoJTIwYWNjZWxlcmF0b3IuYWNjdW11bGF0ZShtb2RlbCklM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjAlRTklQTIlODQlRTYlQjUlOEIlRTUlOTklQUElRTUlQTMlQjAlRTYlQUUlOEIlRTUlQjclQUUlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBub2lzZV9wcmVkJTIwJTNEJTIwbW9kZWwobm9pc3lfaW1hZ2VzJTJDJTIwdGltZXN0ZXBzJTJDJTIwcmV0dXJuX2RpY3QlM0RGYWxzZSklNUIwJTVEJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwbG9zcyUyMCUzRCUyMEYubXNlX2xvc3Mobm9pc2VfcHJlZCUyQyUyMG5vaXNlKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGFjY2VsZXJhdG9yLmJhY2t3YXJkKGxvc3MpJTBBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwaWYlMjBhY2NlbGVyYXRvci5zeW5jX2dyYWRpZW50cyUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGFjY2VsZXJhdG9yLmNsaXBfZ3JhZF9ub3JtXyhtb2RlbC5wYXJhbWV0ZXJzKCklMkMlMjAxLjApJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwb3B0aW1pemVyLnN0ZXAoKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGxyX3NjaGVkdWxlci5zdGVwKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBvcHRpbWl6ZXIuemVyb19ncmFkKCklMEElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBwcm9ncmVzc19iYXIudXBkYXRlKDEpJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwbG9ncyUyMCUzRCUyMCU3QiUyMmxvc3MlMjIlM0ElMjBsb3NzLmRldGFjaCgpLml0ZW0oKSUyQyUyMCUyMmxyJTIyJTNBJTIwbHJfc2NoZWR1bGVyLmdldF9sYXN0X2xyKCklNUIwJTVEJTJDJTIwJTIyc3RlcCUyMiUzQSUyMGdsb2JhbF9zdGVwJTdEJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwcHJvZ3Jlc3NfYmFyLnNldF9wb3N0Zml4KCoqbG9ncyklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBhY2NlbGVyYXRvci5sb2cobG9ncyUyQyUyMHN0ZXAlM0RnbG9iYWxfc3RlcCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBnbG9iYWxfc3RlcCUyMCUyQiUzRCUyMDElMEElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjAlRTYlQUYlOEYlRTQlQjglQUElMjBlcG9jaCUyMCVFNSU5MCU4RSVFNSU4RiVBRiVFNCVCQiVBNSVFOSU4MCU4OSVFNiU4QiVBOSVFNyU5NCVBOCUyMGV2YWx1YXRlKCklMjAlRTklODclODclRTYlQTAlQjclRTQlQjglODAlRTQlQkElOUIlRTYlQkMlOTQlRTclQTQlQkElRTUlOUIlQkUlRTUlODMlOEYlRUYlQkMlOEMlRTUlQjklQjYlRTQlQkYlOUQlRTUlQUQlOTglRTYlQTglQTElRTUlOUUlOEIlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBpZiUyMGFjY2VsZXJhdG9yLmlzX21haW5fcHJvY2VzcyUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHBpcGVsaW5lJTIwJTNEJTIwRERQTVBpcGVsaW5lKHVuZXQlM0RhY2NlbGVyYXRvci51bndyYXBfbW9kZWwobW9kZWwpJTJDJTIwc2NoZWR1bGVyJTNEbm9pc2Vfc2NoZWR1bGVyKSUwQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlmJTIwKGVwb2NoJTIwJTJCJTIwMSklMjAlMjUlMjBjb25maWcuc2F2ZV9pbWFnZV9lcG9jaHMlMjAlM0QlM0QlMjAwJTIwb3IlMjBlcG9jaCUyMCUzRCUzRCUyMGNvbmZpZy5udW1fZXBvY2hzJTIwLSUyMDElM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBldmFsdWF0ZShjb25maWclMkMlMjBlcG9jaCUyQyUyMHBpcGVsaW5lKSUwQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlmJTIwKGVwb2NoJTIwJTJCJTIwMSklMjAlMjUlMjBjb25maWcuc2F2ZV9tb2RlbF9lcG9jaHMlMjAlM0QlM0QlMjAwJTIwb3IlMjBlcG9jaCUyMCUzRCUzRCUyMGNvbmZpZy5udW1fZXBvY2hzJTIwLSUyMDElM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBpZiUyMGNvbmZpZy5wdXNoX3RvX2h1YiUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHVwbG9hZF9mb2xkZXIoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwcmVwb19pZCUzRHJlcG9faWQlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBmb2xkZXJfcGF0aCUzRGNvbmZpZy5vdXRwdXRfZGlyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwY29tbWl0X21lc3NhZ2UlM0RmJTIyRXBvY2glMjAlN0JlcG9jaCU3RCUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlnbm9yZV9wYXR0ZXJucyUzRCU1QiUyMnN0ZXBfKiUyMiUyQyUyMCUyMmVwb2NoXyolMjIlNUQlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjApJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZWxzZSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHBpcGVsaW5lLnNhdmVfcHJldHJhaW5lZChjb25maWcub3V0cHV0X2Rpcik=",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> accelerate <span class="hljs-keyword">import</span> Accelerator
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> huggingface_hub <span class="hljs-keyword">import</span> create_repo, upload_folder
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> tqdm.auto <span class="hljs-keyword">import</span> tqdm
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> pathlib <span class="hljs-keyword">import</span> Path
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> os
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">def</span> <span class="hljs-title function_">train_loop</span>(<span class="hljs-params">config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler</span>):
<span class="hljs-meta">... </span> <span class="hljs-comment"># 初始化 accelerator 和 tensorboard 日志</span>
<span class="hljs-meta">... </span> accelerator = Accelerator(
<span class="hljs-meta">... </span> mixed_precision=config.mixed_precision,
<span class="hljs-meta">... </span> gradient_accumulation_steps=config.gradient_accumulation_steps,
<span class="hljs-meta">... </span> log_with=<span class="hljs-string">&quot;tensorboard&quot;</span>,
<span class="hljs-meta">... </span> project_dir=os.path.join(config.output_dir, <span class="hljs-string">&quot;logs&quot;</span>),
<span class="hljs-meta">... </span> )
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> accelerator.is_main_process:
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> config.output_dir <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-literal">None</span>:
<span class="hljs-meta">... </span> os.makedirs(config.output_dir, exist_ok=<span class="hljs-literal">True</span>)
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> config.push_to_hub:
<span class="hljs-meta">... </span> repo_id = create_repo(
<span class="hljs-meta">... </span> repo_id=config.hub_model_id <span class="hljs-keyword">or</span> Path(config.output_dir).name, exist_ok=<span class="hljs-literal">True</span>
<span class="hljs-meta">... </span> ).repo_id
<span class="hljs-meta">... </span> accelerator.init_trackers(<span class="hljs-string">&quot;train_example&quot;</span>)
<span class="hljs-meta">... </span> <span class="hljs-comment"># 准备所有对象</span>
<span class="hljs-meta">... </span> <span class="hljs-comment"># 不需要记住固定顺序，只要解包时和传给 prepare 的顺序一致即可。</span>
<span class="hljs-meta">... </span> model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
<span class="hljs-meta">... </span> model, optimizer, train_dataloader, lr_scheduler
<span class="hljs-meta">... </span> )
<span class="hljs-meta">... </span> global_step = <span class="hljs-number">0</span>
<span class="hljs-meta">... </span> <span class="hljs-comment"># 开始训练模型</span>
<span class="hljs-meta">... </span> <span class="hljs-keyword">for</span> epoch <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(config.num_epochs):
<span class="hljs-meta">... </span> progress_bar = tqdm(total=<span class="hljs-built_in">len</span>(train_dataloader), disable=<span class="hljs-keyword">not</span> accelerator.is_local_main_process)
<span class="hljs-meta">... </span> progress_bar.set_description(<span class="hljs-string">f&quot;Epoch <span class="hljs-subst">{epoch}</span>&quot;</span>)
<span class="hljs-meta">... </span> <span class="hljs-keyword">for</span> step, batch <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(train_dataloader):
<span class="hljs-meta">... </span> clean_images = batch[<span class="hljs-string">&quot;images&quot;</span>]
<span class="hljs-meta">... </span> <span class="hljs-comment"># 为图像采样噪声</span>
<span class="hljs-meta">... </span> noise = torch.randn(clean_images.shape, device=clean_images.device)
<span class="hljs-meta">... </span> bs = clean_images.shape[<span class="hljs-number">0</span>]
<span class="hljs-meta">... </span> <span class="hljs-comment"># 为每张图像随机采样一个时间步</span>
<span class="hljs-meta">... </span> timesteps = torch.randint(
<span class="hljs-meta">... </span> <span class="hljs-number">0</span>, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,
<span class="hljs-meta">... </span> dtype=torch.int64
<span class="hljs-meta">... </span> )
<span class="hljs-meta">... </span> <span class="hljs-comment"># 按照每个时间步对应的噪声强度给干净图像加噪</span>
<span class="hljs-meta">... </span> <span class="hljs-comment"># （这就是前向扩散过程）</span>
<span class="hljs-meta">... </span> noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
<span class="hljs-meta">... </span> <span class="hljs-keyword">with</span> accelerator.accumulate(model):
<span class="hljs-meta">... </span> <span class="hljs-comment"># 预测噪声残差</span>
<span class="hljs-meta">... </span> noise_pred = model(noisy_images, timesteps, return_dict=<span class="hljs-literal">False</span>)[<span class="hljs-number">0</span>]
<span class="hljs-meta">... </span> loss = F.mse_loss(noise_pred, noise)
<span class="hljs-meta">... </span> accelerator.backward(loss)
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> accelerator.sync_gradients:
<span class="hljs-meta">... </span> accelerator.clip_grad_norm_(model.parameters(), <span class="hljs-number">1.0</span>)
<span class="hljs-meta">... </span> optimizer.step()
<span class="hljs-meta">... </span> lr_scheduler.step()
<span class="hljs-meta">... </span> optimizer.zero_grad()
<span class="hljs-meta">... </span> progress_bar.update(<span class="hljs-number">1</span>)
<span class="hljs-meta">... </span> logs = {<span class="hljs-string">&quot;loss&quot;</span>: loss.detach().item(), <span class="hljs-string">&quot;lr&quot;</span>: lr_scheduler.get_last_lr()[<span class="hljs-number">0</span>], <span class="hljs-string">&quot;step&quot;</span>: global_step}
<span class="hljs-meta">... </span> progress_bar.set_postfix(**logs)
<span class="hljs-meta">... </span> accelerator.log(logs, step=global_step)
<span class="hljs-meta">... </span> global_step += <span class="hljs-number">1</span>
<span class="hljs-meta">... </span> <span class="hljs-comment"># 每个 epoch 后可以选择用 evaluate() 采样一些演示图像，并保存模型</span>
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> accelerator.is_main_process:
<span class="hljs-meta">... </span> pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> (epoch + <span class="hljs-number">1</span>) % config.save_image_epochs == <span class="hljs-number">0</span> <span class="hljs-keyword">or</span> epoch == config.num_epochs - <span class="hljs-number">1</span>:
<span class="hljs-meta">... </span> evaluate(config, epoch, pipeline)
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> (epoch + <span class="hljs-number">1</span>) % config.save_model_epochs == <span class="hljs-number">0</span> <span class="hljs-keyword">or</span> epoch == config.num_epochs - <span class="hljs-number">1</span>:
<span class="hljs-meta">... </span> <span class="hljs-keyword">if</span> config.push_to_hub:
<span class="hljs-meta">... </span> upload_folder(
<span class="hljs-meta">... </span> repo_id=repo_id,
<span class="hljs-meta">... </span> folder_path=config.output_dir,
<span class="hljs-meta">... </span> commit_message=<span class="hljs-string">f&quot;Epoch <span class="hljs-subst">{epoch}</span>&quot;</span>,
<span class="hljs-meta">... </span> ignore_patterns=[<span class="hljs-string">&quot;step_*&quot;</span>, <span class="hljs-string">&quot;epoch_*&quot;</span>],
<span class="hljs-meta">... </span> )
<span class="hljs-meta">... </span> <span class="hljs-keyword">else</span>:
<span class="hljs-meta">... </span> pipeline.save_pretrained(config.output_dir)`,lang:"py",wrap:!1});var R=s(Q,4);l(R,{code:"ZnJvbSUyMGFjY2VsZXJhdGUlMjBpbXBvcnQlMjBub3RlYm9va19sYXVuY2hlciUwQSUwQWFyZ3MlMjAlM0QlMjAoY29uZmlnJTJDJTIwbW9kZWwlMkMlMjBub2lzZV9zY2hlZHVsZXIlMkMlMjBvcHRpbWl6ZXIlMkMlMjB0cmFpbl9kYXRhbG9hZGVyJTJDJTIwbHJfc2NoZWR1bGVyKSUwQSUwQW5vdGVib29rX2xhdW5jaGVyKHRyYWluX2xvb3AlMkMlMjBhcmdzJTJDJTIwbnVtX3Byb2Nlc3NlcyUzRDEp",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> accelerate <span class="hljs-keyword">import</span> notebook_launcher
<span class="hljs-meta">&gt;&gt;&gt; </span>args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)
<span class="hljs-meta">&gt;&gt;&gt; </span>notebook_launcher(train_loop, args, num_processes=<span class="hljs-number">1</span>)`,lang:"py",wrap:!1});var N=s(R,4);l(N,{code:"aW1wb3J0JTIwZ2xvYiUwQSUwQXNhbXBsZV9pbWFnZXMlMjAlM0QlMjBzb3J0ZWQoZ2xvYi5nbG9iKGYlMjIlN0Jjb25maWcub3V0cHV0X2RpciU3RCUyRnNhbXBsZXMlMkYqLnBuZyUyMikpJTBBSW1hZ2Uub3BlbihzYW1wbGVfaW1hZ2VzJTVCLTElNUQp",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> glob
<span class="hljs-meta">&gt;&gt;&gt; </span>sample_images = <span class="hljs-built_in">sorted</span>(glob.glob(<span class="hljs-string">f&quot;<span class="hljs-subst">{config.output_dir}</span>/samples/*.png&quot;</span>))
<span class="hljs-meta">&gt;&gt;&gt; </span>Image.<span class="hljs-built_in">open</span>(sample_images[-<span class="hljs-number">1</span>])`,lang:"py",wrap:!1});var B=s(N,4);a(B,{title:"下一步",local:"下一步",headingTag:"h2"});var _=s(B,6);Y(_,{source:"https://github.com/huggingface/diffusers/blob/main/docs/source/zh/tutorials/basic_training.md"}),H(2),G(F,n),x()}export{ns as component};

Xet Storage Details

Size:
55.9 kB
·
Xet hash:
59533e166a798797fe8fe05adca93f324c6ee46a4ec7f8d6eda2e6040495e54e

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.