"
]
},
"metadata": {},
"execution_count": 23
}
]
},
{
"cell_type": "markdown",
"source": [
"Напишем предобработку данных. Для моделей Diffusers важно, чтобы **пиксели изображения лежали в диапазоне $\\small [-1, 1]$**. Для этого используется `transforms.Normalize` c $\\small \\text{mean}=0.5$, $\\small \\text{std}=0.5$.\n",
"\n",
"*Важно: статистику по датасету для нормализации считать не нужно. Если картинки разных размеров, приведите их к одному размеру.*"
],
"metadata": {
"id": "lF51cSraWwkd"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zQytGLtQue6O"
},
"outputs": [],
"source": [
"from torchvision import transforms\n",
"\n",
"preprocess = transforms.Compose(\n",
" [\n",
" transforms.ToTensor(),\n",
" transforms.Normalize(\n",
" (0.5,), (0.5,)\n",
" ), # it is typical normalization for Diffusers\n",
" ]\n",
")"
]
},
{
"cell_type": "markdown",
"source": [
"Передадим преобразование в датасет, чтобы преобразовывать данные \"на лету\" во время обучения."
],
"metadata": {
"id": "n8M3yOEZbMh8"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tGwf-zSmutiF"
},
"outputs": [],
"source": [
"def transform(examples):\n",
" images = [preprocess(image) for image in examples[\"image\"]]\n",
" return {\"images\": images}\n",
"\n",
"\n",
"dataset.set_transform(transform)"
]
},
{
"cell_type": "markdown",
"source": [
"Создадим загрузчик данных, используя информацию из `config`."
],
"metadata": {
"id": "7Qf3L4uSbQzT"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5dbMoMoyuuVC"
},
"outputs": [],
"source": [
"import torch\n",
"\n",
"train_dataloader = torch.utils.data.DataLoader(\n",
" dataset, batch_size=config.train_batch_size, shuffle=True\n",
")"
]
},
{
"cell_type": "markdown",
"source": [
" ## Создание модели"
],
"metadata": {
"id": "DS6q6ytjrlnM"
}
},
{
"cell_type": "markdown",
"source": [
"Denoising-модель можно легко создать, передав классу модели необходимые параметры:\n",
"- `sample_size` — высота и ширина изображения, тут 28×28.\n",
"- `in_channels`, `out_channels` — количество каналов изображения, тут 1 (черно-белое изображение).\n",
"- `down_block_types` и `up_block_types` — тип блоков понижения и повышения размерности, используемых для создания архитектуры U-Net. В данном случае у нас только по одному блоку с Attention ближе к середине U-Net:\n",
" * `\"DownBlock2D\"` — это стандратный ResNet-блок,\n",
" * `\"AttnDownBlock2D\"` — ResNet-блок с Attention.\n",
"- `block_out_channels` — количество выходных каналов блоков понижения размерности. Используется в обратном порядке для количества входных каналов блоков повышения размерности.\n",
"- `layers_per_block` — количество блоков ResNet, присутствующих в каждом блоке U-Net.\n"
],
"metadata": {
"id": "KW1H5TA8tUi7"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vigoHaktuyfn"
},
"outputs": [],
"source": [
"from diffusers import UNet2DModel\n",
"\n",
"model = UNet2DModel(\n",
" sample_size=config.image_size, # the target image resolution\n",
" in_channels=1, # the number of input channels\n",
" out_channels=1, # the number of output channels\n",
" layers_per_block=2, # how many ResNet layers to use per U-Net block\n",
" block_out_channels=(\n",
" 32,\n",
" 64,\n",
" 128,\n",
" ), # the number of output channels for each U-Net block\n",
" attention_head_dim=8,\n",
" down_block_types=(\n",
" \"DownBlock2D\", # a regular ResNet downsampling block\n",
" \"AttnDownBlock2D\", # a ResNet downsampling block with spatial self-attention\n",
" \"DownBlock2D\",\n",
" ),\n",
" up_block_types=(\n",
" \"UpBlock2D\", # a regular ResNet upsampling block\n",
" \"AttnUpBlock2D\", # a ResNet upsampling block with spatial self-attention\n",
" \"UpBlock2D\",\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"source": [
"Распечатаем архитектуру:"
],
"metadata": {
"id": "hYWnJ0B7mKBw"
}
},
{
"cell_type": "code",
"source": [
"print(model)"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "7rPf6hsQuF0m",
"outputId": "f09b9b06-f0fd-49dc-f463-63d7e3cab394"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"UNet2DModel(\n",
" (conv_in): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_proj): Timesteps()\n",
" (time_embedding): TimestepEmbedding(\n",
" (linear_1): Linear(in_features=32, out_features=128, bias=True)\n",
" (act): SiLU()\n",
" (linear_2): Linear(in_features=128, out_features=128, bias=True)\n",
" )\n",
" (down_blocks): ModuleList(\n",
" (0): DownBlock2D(\n",
" (resnets): ModuleList(\n",
" (0-1): 2 x ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 32, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)\n",
" (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" )\n",
" )\n",
" (downsamplers): ModuleList(\n",
" (0): Downsample2D(\n",
" (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
" )\n",
" )\n",
" )\n",
" (1): AttnDownBlock2D(\n",
" (attentions): ModuleList(\n",
" (0-1): 2 x Attention(\n",
" (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (to_q): Linear(in_features=64, out_features=64, bias=True)\n",
" (to_k): Linear(in_features=64, out_features=64, bias=True)\n",
" (to_v): Linear(in_features=64, out_features=64, bias=True)\n",
" (to_out): ModuleList(\n",
" (0): Linear(in_features=64, out_features=64, bias=True)\n",
" (1): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (resnets): ModuleList(\n",
" (0): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 32, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)\n",
" (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" (1): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)\n",
" (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" )\n",
" )\n",
" (downsamplers): ModuleList(\n",
" (0): Downsample2D(\n",
" (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
" )\n",
" )\n",
" )\n",
" (2): DownBlock2D(\n",
" (resnets): ModuleList(\n",
" (0): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=128, bias=True)\n",
" (norm2): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" (1): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=128, bias=True)\n",
" (norm2): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" )\n",
" )\n",
" )\n",
" )\n",
" (up_blocks): ModuleList(\n",
" (0): UpBlock2D(\n",
" (resnets): ModuleList(\n",
" (0-1): 2 x ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 256, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=128, bias=True)\n",
" (norm2): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" (2): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 192, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(192, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=128, bias=True)\n",
" (norm2): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(192, 128, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" )\n",
" (upsamplers): ModuleList(\n",
" (0): Upsample2D(\n",
" (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" )\n",
" )\n",
" )\n",
" (1): AttnUpBlock2D(\n",
" (attentions): ModuleList(\n",
" (0-2): 3 x Attention(\n",
" (group_norm): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (to_q): Linear(in_features=64, out_features=64, bias=True)\n",
" (to_k): Linear(in_features=64, out_features=64, bias=True)\n",
" (to_v): Linear(in_features=64, out_features=64, bias=True)\n",
" (to_out): ModuleList(\n",
" (0): Linear(in_features=64, out_features=64, bias=True)\n",
" (1): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (resnets): ModuleList(\n",
" (0): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 192, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)\n",
" (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" (1): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)\n",
" (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" (2): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 96, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(96, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=64, bias=True)\n",
" (norm2): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(96, 64, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" )\n",
" (upsamplers): ModuleList(\n",
" (0): Upsample2D(\n",
" (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" )\n",
" )\n",
" )\n",
" (2): UpBlock2D(\n",
" (resnets): ModuleList(\n",
" (0): ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 96, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)\n",
" (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(96, 32, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" (1-2): 2 x ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 64, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=32, bias=True)\n",
" (norm2): GroupNorm(32, 32, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" (conv_shortcut): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
" )\n",
" )\n",
" )\n",
" )\n",
" (mid_block): UNetMidBlock2D(\n",
" (attentions): ModuleList(\n",
" (0): Attention(\n",
" (group_norm): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (to_q): Linear(in_features=128, out_features=128, bias=True)\n",
" (to_k): Linear(in_features=128, out_features=128, bias=True)\n",
" (to_v): Linear(in_features=128, out_features=128, bias=True)\n",
" (to_out): ModuleList(\n",
" (0): Linear(in_features=128, out_features=128, bias=True)\n",
" (1): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (resnets): ModuleList(\n",
" (0-1): 2 x ResnetBlock2D(\n",
" (norm1): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (time_emb_proj): Linear(in_features=128, out_features=128, bias=True)\n",
" (norm2): GroupNorm(32, 128, eps=1e-05, affine=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (nonlinearity): SiLU()\n",
" )\n",
" )\n",
" )\n",
" (conv_norm_out): GroupNorm(32, 32, eps=1e-05, affine=True)\n",
" (conv_act): SiLU()\n",
" (conv_out): Conv2d(32, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
")\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"Проверим работоспособность модели и соответствие выхода модели входу:"
],
"metadata": {
"id": "aZLXQTyRmOBd"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "cK-kSDE-vgrg",
"outputId": "29583e55-827d-43cd-e5a0-afe827afe697"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Input shape: torch.Size([1, 1, 28, 28])\n"
]
}
],
"source": [
"sample_image = dataset[0][\"images\"].unsqueeze(0)\n",
"print(\"Input shape:\", sample_image.shape)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "S_qqVYMbvg9K",
"outputId": "33ca3156-16c7-467f-cd7f-8a5a605629a2"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Output shape: torch.Size([1, 1, 28, 28])\n"
]
}
],
"source": [
"print(\"Output shape:\", model(sample_image, timestep=0).sample.shape)"
]
},
{
"cell_type": "markdown",
"source": [
"## Создание планировщика"
],
"metadata": {
"id": "JLymqwkAmwSJ"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Qi28TpM9wEK7"
},
"outputs": [],
"source": [
"noise_scheduler = DDPMScheduler(num_train_timesteps=1000)"
]
},
{
"cell_type": "markdown",
"source": [
"Для $\\large t$-го шага прямого диффузионного процесса зашумленное изображение $\\large \\mathbf{x}_t$ выражается через исходное изображение $\\large \\mathbf{x}_0$ следующим образом:\n",
"\n",
"$$ \\large \\mathbf{x}_t = \\sqrt{\\overline{\\alpha}_t}\\mathbf{x}_0 + \\sqrt{1 - \\overline{\\alpha}_t}\\bar{\\mathbf{n}}_{0}, $$\n",
"\n",
"где $\\large \\bar{\\alpha}_t = \\prod_{s=1}^t \\alpha_s$, а $\\large \\bar{\\mathbf{n}}_{0} \\sim \\mathcal{N}(\\mathbf{0}, \\mathbf{I})$.\n",
"\n",
"Посмотрим, какие значения выдает планировщик:"
],
"metadata": {
"id": "bpVH_RiuubBU"
}
},
{
"cell_type": "code",
"source": [
"plt.plot(\n",
" noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r\"${\\sqrt{\\bar{\\alpha}_t}}$\"\n",
")\n",
"plt.plot(\n",
" (1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5,\n",
" label=r\"$\\sqrt{(1 - \\bar{\\alpha}_t)}$\",\n",
")\n",
"plt.legend(fontsize=\"x-large\")\n",
"plt.plot()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 447
},
"id": "k205_Q4Wnb8d",
"outputId": "f3a6fa69-d433-4c50-aa97-86fb5ab29a0f"
},
"execution_count": null,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"[]"
]
},
"metadata": {},
"execution_count": 32
},
{
"output_type": "display_data",
"data": {
"text/plain": [
"![](\"https://edunet.kea.su/repo/EduNet-content/dev-2.0/msu_ai.png\")