{"id":664,"date":"2024-04-17T00:18:29","date_gmt":"2024-04-16T16:18:29","guid":{"rendered":"http:\/\/tobykskgd.life\/?p=664"},"modified":"2024-11-14T22:18:38","modified_gmt":"2024-11-14T14:18:38","slug":"28","status":"publish","type":"post","link":"https:\/\/tobykskgd.life\/index.php\/28\/","title":{"rendered":"\u674e\u5b8f\u6bc5\u673a\u5668\u5b66\u4e60\u8bfe\u7a0b\u7b14\u8bb0EP18"},"content":{"rendered":"\n

\u3010HW6\u3011Diffusion Model0.1\u674e\u5b8f\u6bc52021\/2022\u6625\u673a\u5668\u5b66\u4e60\u8bfe\u7a0b\u7b14\u8bb0EP18(P68-P70\uff1f)<\/p>\n\n\n\n

\"\"<\/div><\/figure>\n\n\n\n

\u4ece\u4eca\u5929\u5f00\u59cb\u6211\u5c06\u5b66\u4e60\u674e\u5b8f\u6bc5\u6559\u6388\u7684\u673a\u5668\u5b66\u4e60\u89c6\u9891\uff0c\u4e0b\u9762\u662f\u8bfe\u7a0b\u7684\u8fde\u63a5(\u5f3a\u63a8)\u674e\u5b8f\u6bc52021\/2022\u6625\u673a\u5668\u5b66\u4e60\u8bfe\u7a0b_\u54d4\u54e9\u54d4\u54e9_bilibili<\/a>\u3002\u4e00\u5171\u6709155\u4e2a\u89c6\u9891\uff0c\u4e89\u53d6\u90fd\u5b66\u4e60\u5b8c\u6210\u5427\u3002<\/p>\n\n\n\n

\u90a3\u4e48\u9996\u5148\u8fd9\u95e8\u8bfe\u7a0b\u9700\u8981\u6709\u4e00\u5b9a\u7684\u4ee3\u7801\u57fa\u7840\uff0c\u7b80\u5355\u5b66\u4e60\u4e00\u4e0bPython\u7684\u57fa\u672c\u7528\u6cd5\uff0c\u8fd8\u6709\u91cc\u9762\u7684NumPy\u5e93\u7b49\u7b49\u7684\u57fa\u672c\u77e5\u8bc6\u3002\u518d\u5c31\u662f\u6570\u5b66\u65b9\u9762\u7684\u57fa\u7840\u5566\uff0c\u5fae\u79ef\u5206\u3001\u7ebf\u6027\u4ee3\u6570\u548c\u6982\u7387\u8bba\u7684\u57fa\u7840\u90fd\u662f\u542c\u61c2\u8fd9\u95e8\u8bfe\u5fc5\u987b\u7684\u3002<\/p>\n\n\n\n

\n\n\n\n

\u6807\u9898\u8fd9\u91cc\u6211\u5728\u5bf9\u5e94\u7684\u8282\u6570\uff08P68-P70\uff1f\uff09\u4e0a\u9762\u6253\u4e86\u95ee\u53f7\uff0c\u56e0\u4e3a\u539f\u672c2021\/2022\u7684\u8bfe\u7a0b\u8fd9\u91cc\u7684HW6\u505a\u7684\u662f\u7528GAN\u53bb\u751f\u6210\u56fe\u50cf\uff0c\u4e0d\u8fc7\u6211\u8fd0\u884c\u7684\u662f2023\u7684\u4f5c\u4e1a\uff0c\u7136\u540e2023\u5e74\u7528\u7684\u662fDiffusion Model\u53bb\u751f\u6210\u6a21\u578b\uff0c\u6240\u4ee5\u6807\u9898\u4e0a\u9762\u7684\u7ae0\u8282\u5176\u5b9e\u662f\u4e0d\u5bf9\u5e94\u7684\u3002\u7136\u540e\u6211\u4e5f\u662f\u628a2023\u5e74\u6709\u5173\u4e8e\u56fe\u50cf\u751f\u6210\u7684\u8865\u5145\u63d0\u5230\u524d\u9762\u6765\u5b66\u4e86\uff0c\u7b14\u8bb0\u4e5f\u662f\u5728\u65b0\u76842023\u5e74\u5173\u4e8e\u674e\u5b8f\u6bc5\u6559\u6388\u7684\u8bfe\u7a0b\u91cc\u9762\u3002<\/p>\n\n\n\n

\u8fd9\u91cc\u987a\u4fbf\u8981\u8bf4\u7684\u662f\uff0c\u56e0\u4e3a\u8fd9\u6b21\u4f5c\u4e1a\u4e0a\u7684\u6570\u636e\u96c6\u5728kaggle\u4e0a\u9762\uff0c\u6211\u8bbf\u95ee\u4e0d\u4e0a\uff0c\u4e8e\u662f\u8bad\u7ec3\u7528\u7684\u6570\u636e\u96c6\u662f\u6211\u81ea\u5df1\u4e0a\u4f20\u7684\uff0c\u4f46\u662f\u56e0\u4e3a\u539f\u672c\u4f5c\u4e1a\u7684\u6570\u636e\u96c6\u6709\u4e03\u4e07\u591a\u5f20\u7684\u56fe\u7247\uff0ccolab\u4e00\u4e0a\u4f20\u591a\u7684\u8d44\u6599\u5c31\u4f1a\u5d29\u6389\uff0c\u6240\u4ee5\u8fd9\u91cc\u6211\u53ea\u4e0a\u4f20\u4e86\u4e03\u767e\u591a\u5f20\u56fe\u7247\u4f5c\u4e3a\u8bad\u7ec3\u8d44\u6599\u96c6\uff0c\u7136\u540e\u8fd9\u4e2a\u4ee3\u7801\u8dd1\u7684\u65f6\u95f4\u4e5f\u5f88\u957f\uff0c66%\u5c31\u7528\u4e867\u4e2a\u5c0f\u65f6\u5de6\u53f3\uff0c\u6240\u4ee5\u6211\u5c31\u572866%\u8fd9\u91cc\u5148\u505c\u4e0b\u6765\u4e86\uff0c\u4e0b\u9762\u8fd9\u5f20\u56fe\u7247\u5c31\u662f\u8dd1\u4e866\u8f6e\u4e4b\u540e\u751f\u6210\u7684\u56fe\u50cf\uff1a<\/p>\n\n\n\n

\"\"<\/div><\/figure>\n\n\n\n
\n\n\n\n

\u4e0b\u9762\u662f\u4f5c\u4e1a\u7684\u4ee3\u7801\uff1a<\/p>\n\n\n\n

expand_less<\/p>\n\n\n\n

\n\n\n\n

Import Packages and Set Seeds<\/h1>\n\n\n\n
!pip install einops\n!pip install transformers\n!pip install ema_pytorch\n!pip install accelerate<\/code><\/pre>\n\n\n\n
Collecting einops Downloading einops-0.7.0-py3-none-any.whl (44 kB) \u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501 44.6\/44.6 kB 1.2 MB\/s eta 0:00:00 Installing collected packages: einops Successfully installed einops-0.7.0 Requirement already satisfied: transformers in \/usr\/local\/lib\/python3.10\/dist-packages (4.38.2) Requirement already satisfied: filelock in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (3.13.4) Requirement already satisfied: huggingface-hub<1.0,>=0.19.3 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (0.20.3) Requirement already satisfied: numpy>=1.17 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (1.25.2) Requirement already satisfied: packaging>=20.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (24.0) Requirement already satisfied: pyyaml>=5.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (6.0.1) Requirement already satisfied: regex!=2019.12.17 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (2023.12.25) Requirement already satisfied: requests in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (2.31.0) Requirement already satisfied: tokenizers<0.19,>=0.14 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (0.15.2) Requirement already satisfied: safetensors>=0.4.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (0.4.2) Requirement already satisfied: tqdm>=4.27 in \/usr\/local\/lib\/python3.10\/dist-packages (from transformers) (4.66.2) Requirement already satisfied: fsspec>=2023.5.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from huggingface-hub<1.0,>=0.19.3->transformers) (2023.6.0) Requirement already satisfied: typing-extensions>=3.7.4.3 in \/usr\/local\/lib\/python3.10\/dist-packages (from huggingface-hub<1.0,>=0.19.3->transformers) (4.11.0) Requirement already satisfied: charset-normalizer<4,>=2 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->transformers) (3.3.2) Requirement already satisfied: idna<4,>=2.5 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->transformers) (3.6) Requirement already satisfied: urllib3<3,>=1.21.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->transformers) (2.0.7) Requirement already satisfied: certifi>=2017.4.17 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->transformers) (2024.2.2) Collecting ema_pytorch Downloading ema_pytorch-0.4.5-py3-none-any.whl (8.4 kB) Collecting beartype (from ema_pytorch) Downloading beartype-0.18.2-py3-none-any.whl (903 kB) \u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501 903.7\/903.7 kB 10.0 MB\/s eta 0:00:00 Requirement already satisfied: torch>=1.6 in \/usr\/local\/lib\/python3.10\/dist-packages (from ema_pytorch) (2.2.1+cu121) Requirement already satisfied: filelock in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (3.13.4) Requirement already satisfied: typing-extensions>=4.8.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (4.11.0) Requirement already satisfied: sympy in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (1.12) Requirement already satisfied: networkx in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (3.3) Requirement already satisfied: jinja2 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (3.1.3) Requirement already satisfied: fsspec in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (2023.6.0) Collecting nvidia-cuda-nvrtc-cu12==12.1.105 (from torch>=1.6->ema_pytorch) Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (23.7 MB) Collecting nvidia-cuda-runtime-cu12==12.1.105 (from torch>=1.6->ema_pytorch) Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (823 kB) Collecting nvidia-cuda-cupti-cu12==12.1.105 (from torch>=1.6->ema_pytorch) Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (14.1 MB) Collecting nvidia-cudnn-cu12==8.9.2.26 (from torch>=1.6->ema_pytorch) Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl (731.7 MB) Collecting nvidia-cublas-cu12==12.1.3.1 (from torch>=1.6->ema_pytorch) Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl (410.6 MB) Collecting nvidia-cufft-cu12==11.0.2.54 (from torch>=1.6->ema_pytorch) Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl (121.6 MB) Collecting nvidia-curand-cu12==10.3.2.106 (from torch>=1.6->ema_pytorch) Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl (56.5 MB) Collecting nvidia-cusolver-cu12==11.4.5.107 (from torch>=1.6->ema_pytorch) Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl (124.2 MB) Collecting nvidia-cusparse-cu12==12.1.0.106 (from torch>=1.6->ema_pytorch) Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl (196.0 MB) Collecting nvidia-nccl-cu12==2.19.3 (from torch>=1.6->ema_pytorch) Using cached nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl (166.0 MB) Collecting nvidia-nvtx-cu12==12.1.105 (from torch>=1.6->ema_pytorch) Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (99 kB) Requirement already satisfied: triton==2.2.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.6->ema_pytorch) (2.2.0) Collecting nvidia-nvjitlink-cu12 (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.6->ema_pytorch) Using cached nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB) Requirement already satisfied: MarkupSafe>=2.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from jinja2->torch>=1.6->ema_pytorch) (2.1.5) Requirement already satisfied: mpmath>=0.19 in \/usr\/local\/lib\/python3.10\/dist-packages (from sympy->torch>=1.6->ema_pytorch) (1.3.0) Installing collected packages: nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, beartype, nvidia-cusparse-cu12, nvidia-cudnn-cu12, nvidia-cusolver-cu12, ema_pytorch Successfully installed beartype-0.18.2 ema_pytorch-0.4.5 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.19.3 nvidia-nvjitlink-cu12-12.4.127 nvidia-nvtx-cu12-12.1.105 Collecting accelerate Downloading accelerate-0.29.2-py3-none-any.whl (297 kB) \u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501 297.4\/297.4 kB 6.3 MB\/s eta 0:00:00 Requirement already satisfied: numpy>=1.17 in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (1.25.2) Requirement already satisfied: packaging>=20.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (24.0) Requirement already satisfied: psutil in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (5.9.5) Requirement already satisfied: pyyaml in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (6.0.1) Requirement already satisfied: torch>=1.10.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (2.2.1+cu121) Requirement already satisfied: huggingface-hub in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (0.20.3) Requirement already satisfied: safetensors>=0.3.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from accelerate) (0.4.2) Requirement already satisfied: filelock in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (3.13.4) Requirement already satisfied: typing-extensions>=4.8.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (4.11.0) Requirement already satisfied: sympy in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (1.12) Requirement already satisfied: networkx in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (3.3) Requirement already satisfied: jinja2 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (3.1.3) Requirement already satisfied: fsspec in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (2023.6.0) Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (12.1.105) Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (12.1.105) Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (12.1.105) Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (8.9.2.26) Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (12.1.3.1) Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (11.0.2.54) Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (10.3.2.106) Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (11.4.5.107) Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (12.1.0.106) Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (2.19.3) Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (12.1.105) Requirement already satisfied: triton==2.2.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from torch>=1.10.0->accelerate) (2.2.0) Requirement already satisfied: nvidia-nvjitlink-cu12 in \/usr\/local\/lib\/python3.10\/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.10.0->accelerate) (12.4.127) Requirement already satisfied: requests in \/usr\/local\/lib\/python3.10\/dist-packages (from huggingface-hub->accelerate) (2.31.0) Requirement already satisfied: tqdm>=4.42.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from huggingface-hub->accelerate) (4.66.2) Requirement already satisfied: MarkupSafe>=2.0 in \/usr\/local\/lib\/python3.10\/dist-packages (from jinja2->torch>=1.10.0->accelerate) (2.1.5) Requirement already satisfied: charset-normalizer<4,>=2 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->huggingface-hub->accelerate) (3.3.2) Requirement already satisfied: idna<4,>=2.5 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->huggingface-hub->accelerate) (3.6) Requirement already satisfied: urllib3<3,>=1.21.1 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->huggingface-hub->accelerate) (2.0.7) Requirement already satisfied: certifi>=2017.4.17 in \/usr\/local\/lib\/python3.10\/dist-packages (from requests->huggingface-hub->accelerate) (2024.2.2) Requirement already satisfied: mpmath>=0.19 in \/usr\/local\/lib\/python3.10\/dist-packages (from sympy->torch>=1.10.0->accelerate) (1.3.0) Installing collected packages: accelerate Successfully installed accelerate-0.29.2<\/p>\n\n\n\n
import math\nimport copy\nfrom pathlib import Path\nfrom random import random\nfrom functools import partial\nfrom collections import namedtuple\nfrom multiprocessing import cpu_count\n\nimport torch\nfrom torch import nn, einsum\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom torch.optim import Adam\n\nimport torchvision\nfrom torchvision import transforms as T, utils\n\nfrom einops import rearrange, reduce, repeat\nfrom einops.layers.torch import Rearrange\n\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nfrom ema_pytorch import EMA\n\nfrom accelerate import Accelerator\nimport matplotlib.pyplot as plt\nimport os\n\ntorch.backends.cudnn.benchmark = True\ntorch.manual_seed(4096)\n\nif torch.cuda.is_available():\n  torch.cuda.manual_seed(4096)<\/code><\/pre>\n\n\n\n
Step 1: Forward process (Noise scheduler)<\/p>\n\n\n\n
def linear_beta_schedule(timesteps):\n    \"\"\"\n    linear schedule, proposed in original ddpm paper\n    \"\"\"\n    scale = 1000 \/ timesteps\n    beta_start = scale * 0.0001\n    beta_end = scale * 0.02\n    return torch.linspace(beta_start, beta_end, timesteps, dtype = torch.float64)\n\ndef extract(a, t, x_shape):\n    b, *_ = t.shape\n    out = a.gather(-1, t)\n    return out.reshape(b, *((1,) * (len(x_shape) - 1)))<\/code><\/pre>\n\n\n\n
Create dataset<\/p>\n\n\n\n
class Dataset(Dataset):\n    def __init__(\n        self,\n        folder,\n        image_size\n    ):\n        self.folder = folder\n        self.image_size = image_size\n        self.paths = [p for p in Path(f'{folder}').glob(f'**\/*.jpg')]\n        #################################\n        ## TODO: Data Augmentation ##\n        #################################\n        self.transform = T.Compose([\n            T.Resize(image_size),\n            T.ToTensor()\n        ])\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, index):\n        path = self.paths[index]\n        img = Image.open(path)\n        return self.transform(img)<\/code><\/pre>\n\n\n\n
Step 2: The backward process = U-Net<\/p>\n\n\n\n
def exists(x):\n    return x is not None\n\ndef default(val, d):\n    if exists(val):\n        return val\n    return d() if callable(d) else d\n\ndef identity(t, *args, **kwargs):\n    return t\n\ndef cycle(dl):\n    while True:\n        for data in dl:\n            yield data\n\ndef has_int_squareroot(num):\n    return (math.sqrt(num) ** 2) == num\n\ndef num_to_groups(num, divisor):\n    groups = num \/\/ divisor\n    remainder = num % divisor\n    arr = [divisor] * groups\n    if remainder > 0:\n        arr.append(remainder)\n    return arr\n\n# normalization functions\n\ndef normalize_to_neg_one_to_one(img):\n    return img * 2 - 1\n\ndef unnormalize_to_zero_to_one(t):\n    return (t + 1) * 0.5\n\n# small helper modules\n\nclass Residual(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n\n    def forward(self, x, *args, **kwargs):\n        return self.fn(x, *args, **kwargs) + x\n\ndef Upsample(dim, dim_out = None):\n    return nn.Sequential(\n        nn.Upsample(scale_factor = 2, mode = 'nearest'),\n        nn.Conv2d(dim, default(dim_out, dim), 3, padding = 1)\n    )\n\ndef Downsample(dim, dim_out = None):\n    return nn.Sequential(\n        Rearrange('b c (h p1) (w p2) -> b (c p1 p2) h w', p1 = 2, p2 = 2),\n        nn.Conv2d(dim * 4, default(dim_out, dim), 1)\n    )\n\nclass WeightStandardizedConv2d(nn.Conv2d):\n    \"\"\"\n    https://arxiv.org\/abs\/1903.10520\n    weight standardization purportedly works synergistically with group normalization\n    \"\"\"\n    def forward(self, x):\n        eps = 1e-5 if x.dtype == torch.float32 else 1e-3\n\n        weight = self.weight\n        mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')\n        var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))\n        normalized_weight = (weight - mean) * (var + eps).rsqrt()\n\n        return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)\n\nclass LayerNorm(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))\n\n    def forward(self, x):\n        eps = 1e-5 if x.dtype == torch.float32 else 1e-3\n        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)\n        mean = torch.mean(x, dim = 1, keepdim = True)\n        return (x - mean) * (var + eps).rsqrt() * self.g\n\nclass PreNorm(nn.Module):\n    def __init__(self, dim, fn):\n        super().__init__()\n        self.fn = fn\n        self.norm = LayerNorm(dim)\n\n    def forward(self, x):\n        x = self.norm(x)\n        return self.fn(x)\n\n# sinusoidal positional embeds\n\nclass SinusoidalPosEmb(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.dim = dim\n\n    def forward(self, x):\n        device = x.device\n        half_dim = self.dim \/\/ 2\n        emb = math.log(10000) \/ (half_dim - 1)\n        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)\n        emb = x[:, None] * emb[None, :]\n        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n        return emb\n\nclass RandomOrLearnedSinusoidalPosEmb(nn.Module):\n    \"\"\" following @crowsonkb 's lead with random (learned optional) sinusoidal pos emb \"\"\"\n    \"\"\" https:\/\/github.com\/crowsonkb\/v-diffusion-jax\/blob\/master\/diffusion\/models\/danbooru_128.py#L8 \"\"\"\n\n    def __init__(self, dim, is_random = False):\n        super().__init__()\n        assert (dim % 2) == 0\n        half_dim = dim \/\/ 2\n        self.weights = nn.Parameter(torch.randn(half_dim), requires_grad = not is_random)\n\n    def forward(self, x):\n        x = rearrange(x, 'b -> b 1')\n        freqs = x * rearrange(self.weights, 'd -> 1 d') * 2 * math.pi\n        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim = -1)\n        fouriered = torch.cat((x, fouriered), dim = -1)\n        return fouriered\n\n# building block modules\n\nclass Block(nn.Module):\n    def __init__(self, dim, dim_out, groups = 8):\n        super().__init__()\n        self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding = 1)\n        self.norm = nn.GroupNorm(groups, dim_out)\n        self.act = nn.SiLU()\n\n    def forward(self, x, scale_shift = None):\n        x = self.proj(x)\n        x = self.norm(x)\n\n        if exists(scale_shift):\n            scale, shift = scale_shift\n            x = x * (scale + 1) + shift\n\n        x = self.act(x)\n        return x\n\nclass ResnetBlock(nn.Module):\n    def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8):\n        super().__init__()\n        self.mlp = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(time_emb_dim, dim_out * 2)\n        ) if exists(time_emb_dim) else None\n\n        self.block1 = Block(dim, dim_out, groups = groups)\n        self.block2 = Block(dim_out, dim_out, groups = groups)\n        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()\n\n    def forward(self, x, time_emb = None):\n\n        scale_shift = None\n        if exists(self.mlp) and exists(time_emb):\n            time_emb = self.mlp(time_emb)\n            time_emb = rearrange(time_emb, 'b c -> b c 1 1')\n            scale_shift = time_emb.chunk(2, dim = 1)\n\n        h = self.block1(x, scale_shift = scale_shift)\n\n        h = self.block2(h)\n\n        return h + self.res_conv(x)\n\nclass LinearAttention(nn.Module):\n    def __init__(self, dim, heads = 4, dim_head = 32):\n        super().__init__()\n        self.scale = dim_head ** -0.5\n        self.heads = heads\n        hidden_dim = dim_head * heads\n        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)\n\n        self.to_out = nn.Sequential(\n            nn.Conv2d(hidden_dim, dim, 1),\n            LayerNorm(dim)\n        )\n\n    def forward(self, x):\n        b, c, h, w = x.shape\n        qkv = self.to_qkv(x).chunk(3, dim = 1)\n        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)\n\n        q = q.softmax(dim = -2)\n        k = k.softmax(dim = -1)\n\n        q = q * self.scale\n        v = v \/ (h * w)\n\n        context = torch.einsum('b h d n, b h e n -> b h d e', k, v)\n\n        out = torch.einsum('b h d e, b h d n -> b h e n', context, q)\n        out = rearrange(out, 'b h c (x y) -> b (h c) x y', h = self.heads, x = h, y = w)\n        return self.to_out(out)\n\nclass Attention(nn.Module):\n    def __init__(self, dim, heads = 4, dim_head = 32):\n        super().__init__()\n        self.scale = dim_head ** -0.5\n        self.heads = heads\n        hidden_dim = dim_head * heads\n\n        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)\n        self.to_out = nn.Conv2d(hidden_dim, dim, 1)\n\n    def forward(self, x):\n        b, c, h, w = x.shape\n        qkv = self.to_qkv(x).chunk(3, dim = 1)\n        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)\n\n        q = q * self.scale\n\n        sim = torch.einsum('b h d i, b h d j -> b h i j', q, k)\n        attn = sim.softmax(dim = -1)\n        out = torch.einsum('b h i j, b h d j -> b h i d', attn, v)\n\n        out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)\n        return self.to_out(out)\n\n# model\n\nclass Unet(nn.Module):\n    def __init__(\n        self,\n        dim,\n        init_dim = None,\n        out_dim = None,\n        dim_mults=(1, 2, 4, 8),\n        channels = 3,\n        resnet_block_groups = 8,\n        learned_sinusoidal_cond = False,\n        random_fourier_features = False,\n        learned_sinusoidal_dim = 16\n    ):\n        super().__init__()\n\n        # determine dimensions\n\n        self.channels = channels\n\n        init_dim = default(init_dim, dim)\n        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding = 3)\n\n        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]\n        in_out = list(zip(dims[:-1], dims[1:]))\n\n        block_klass = partial(ResnetBlock, groups = resnet_block_groups)\n\n        # time embeddings\n\n        time_dim = dim * 4\n\n        self.random_or_learned_sinusoidal_cond = learned_sinusoidal_cond or random_fourier_features\n\n        if self.random_or_learned_sinusoidal_cond:\n            sinu_pos_emb = RandomOrLearnedSinusoidalPosEmb(learned_sinusoidal_dim, random_fourier_features)\n            fourier_dim = learned_sinusoidal_dim + 1\n        else:\n            sinu_pos_emb = SinusoidalPosEmb(dim)\n            fourier_dim = dim\n\n        self.time_mlp = nn.Sequential(\n            sinu_pos_emb,\n            nn.Linear(fourier_dim, time_dim),\n            nn.GELU(),\n            nn.Linear(time_dim, time_dim)\n        )\n\n        # layers\n\n        self.downs = nn.ModuleList([])\n        self.ups = nn.ModuleList([])\n        num_resolutions = len(in_out)\n\n        for ind, (dim_in, dim_out) in enumerate(in_out):\n            is_last = ind >= (num_resolutions - 1)\n\n            self.downs.append(nn.ModuleList([\n                block_klass(dim_in, dim_in, time_emb_dim = time_dim),\n                block_klass(dim_in, dim_in, time_emb_dim = time_dim),\n                Residual(PreNorm(dim_in, LinearAttention(dim_in))),\n                Downsample(dim_in, dim_out) if not is_last else nn.Conv2d(dim_in, dim_out, 3, padding = 1)\n            ]))\n\n        mid_dim = dims[-1]\n        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)\n        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))\n        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)\n\n        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):\n            is_last = ind == (len(in_out) - 1)\n\n            self.ups.append(nn.ModuleList([\n                block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),\n                block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),\n                Residual(PreNorm(dim_out, LinearAttention(dim_out))),\n                Upsample(dim_out, dim_in) if not is_last else  nn.Conv2d(dim_out, dim_in, 3, padding = 1)\n            ]))\n\n        self.out_dim = default(out_dim, channels)\n\n        self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim)\n        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)\n\n    def forward(self, x, time):\n        x = self.init_conv(x)\n        r = x.clone()\n\n        t = self.time_mlp(time)\n\n        h = []\n\n        for block1, block2, attn, downsample in self.downs:\n            x = block1(x, t)\n            h.append(x)\n\n            x = block2(x, t)\n            x = attn(x)\n            h.append(x)\n\n            x = downsample(x)\n\n        x = self.mid_block1(x, t)\n        x = self.mid_attn(x)\n        x = self.mid_block2(x, t)\n\n        for block1, block2, attn, upsample in self.ups:\n            x = torch.cat((x, h.pop()), dim = 1)\n            x = block1(x, t)\n\n            x = torch.cat((x, h.pop()), dim = 1)\n            x = block2(x, t)\n            x = attn(x)\n\n            x = upsample(x)\n\n        x = torch.cat((x, r), dim = 1)\n\n        x = self.final_res_block(x, t)\n        return self.final_conv(x)\nmodel = Unet(64)<\/code><\/pre>\n\n\n\n
Step 3: The Diffusion Process<\/p>\n\n\n\n
Define diffusion process, including generating noisy models, sample…<\/p>\n\n\n\n
class GaussianDiffusion(nn.Module):\n    def __init__(\n        self,\n        model,\n        *,\n        image_size,\n        timesteps = 1000,\n        beta_schedule = 'linear',\n        auto_normalize = True\n    ):\n        super().__init__()\n        assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)\n        assert not model.random_or_learned_sinusoidal_cond\n\n        self.model = model\n\n        self.channels = self.model.channels\n\n        self.image_size = image_size\n\n\n        if beta_schedule == 'linear':\n            beta_schedule_fn = linear_beta_schedule\n        else:\n            raise ValueError(f'unknown beta schedule {beta_schedule}')\n\n        # calculate beta and other precalculated parameters\n        betas = beta_schedule_fn(timesteps)\n\n        alphas = 1. - betas\n        alphas_cumprod = torch.cumprod(alphas, dim=0)\n        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)\n\n        timesteps, = betas.shape\n        self.num_timesteps = int(timesteps)\n\n        # sampling related parameters\n\n        self.sampling_timesteps = timesteps # default num sampling timesteps to number of timesteps at training\n\n        # helper function to register buffer from float64 to float32\n\n        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))\n\n        register_buffer('betas', betas)\n        register_buffer('alphas_cumprod', alphas_cumprod)\n        register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)\n\n        # calculations for diffusion q(x_t | x_{t-1}) and others\n\n        register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))\n        register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))\n        register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))\n        register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. \/ alphas_cumprod))\n        register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. \/ alphas_cumprod - 1))\n\n        # calculations for posterior q(x_{t-1} | x_t, x_0)\n\n        posterior_variance = betas * (1. - alphas_cumprod_prev) \/ (1. - alphas_cumprod)\n\n        # above: equal to 1. \/ (1. \/ (1. - alpha_cumprod_tm1) + alpha_t \/ beta_t)\n\n        register_buffer('posterior_variance', posterior_variance)\n\n        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain\n\n        register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20)))\n        register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) \/ (1. - alphas_cumprod))\n        register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) \/ (1. - alphas_cumprod))\n\n        # derive loss weight\n        # snr - signal noise ratio\n\n        snr = alphas_cumprod \/ (1 - alphas_cumprod)\n\n        # https:\/\/arxiv.org\/abs\/2303.09556\n\n        maybe_clipped_snr = snr.clone()\n\n        register_buffer('loss_weight', maybe_clipped_snr \/ snr)\n\n        # auto-normalization of data [0, 1] -> [-1, 1] - can turn off by setting it to be False\n\n        self.normalize = normalize_to_neg_one_to_one if auto_normalize else identity\n        self.unnormalize = unnormalize_to_zero_to_one if auto_normalize else identity\n\n    def predict_start_from_noise(self, x_t, t, noise):\n        return (\n            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -\n            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise\n        )\n\n    def predict_noise_from_start(self, x_t, t, x0):\n        return (\n            (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) \/ \\\n            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)\n        )\n\n\n    def q_posterior(self, x_start, x_t, t):\n        posterior_mean = (\n            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +\n            extract(self.posterior_mean_coef2, t, x_t.shape) * x_t\n        )\n        posterior_variance = extract(self.posterior_variance, t, x_t.shape)\n        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)\n        return posterior_mean, posterior_variance, posterior_log_variance_clipped\n\n    def model_predictions(self, x, t, clip_x_start = False, rederive_pred_noise = False):\n        model_output = self.model(x, t)\n        maybe_clip = partial(torch.clamp, min = -1., max = 1.) if clip_x_start else identity\n\n        pred_noise = model_output\n        x_start = self.predict_start_from_noise(x, t, pred_noise)\n        x_start = maybe_clip(x_start)\n\n        if clip_x_start and rederive_pred_noise:\n            pred_noise = self.predict_noise_from_start(x, t, x_start)\n\n        return pred_noise, x_start\n\n    def p_mean_variance(self, x, t, clip_denoised = True):\n        noise, x_start = self.model_predictions(x, t)\n\n        if clip_denoised:\n            x_start.clamp_(-1., 1.)\n\n        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)\n        return model_mean, posterior_variance, posterior_log_variance, x_start\n\n    @torch.no_grad()\n    def p_sample(self, x, t: int):\n        b, *_, device = *x.shape, x.device\n        batched_times = torch.full((b,), t, device = x.device, dtype = torch.long)\n        model_mean, _, model_log_variance, x_start = self.p_mean_variance(x = x, t = batched_times, clip_denoised = True)\n        noise = torch.randn_like(x) if t > 0 else 0. # no noise if t == 0\n        pred_img = model_mean + (0.5 * model_log_variance).exp() * noise\n        return pred_img, x_start\n\n    @torch.no_grad()\n    def p_sample_loop(self, shape, return_all_timesteps = False):\n        batch, device = shape[0], self.betas.device\n\n        img = torch.randn(shape, device = device)\n        imgs = [img]\n\n        x_start = None\n\n        ###########################################\n        ## TODO: plot the sampling process ##\n        ###########################################\n        for t in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step', total = self.num_timesteps):\n            img, x_start = self.p_sample(img, t)\n            imgs.append(img)\n\n        ret = img if not return_all_timesteps else torch.stack(imgs, dim = 1)\n\n        ret = self.unnormalize(ret)\n        return ret\n\n    @torch.no_grad()\n    def sample(self, batch_size = 16, return_all_timesteps = False):\n        image_size, channels = self.image_size, self.channels\n        sample_fn = self.p_sample_loop\n        return sample_fn((batch_size, channels, image_size, image_size), return_all_timesteps = return_all_timesteps)\n\n\n    def q_sample(self, x_start, t, noise=None):\n        noise = default(noise, lambda: torch.randn_like(x_start))\n\n        return (\n            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +\n            extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise\n        )\n\n    @property\n    def loss_fn(self):\n        return F.mse_loss\n\n\n    def p_losses(self, x_start, t, noise = None):\n        b, c, h, w = x_start.shape\n        noise = default(noise, lambda: torch.randn_like(x_start))\n\n        # noise sample\n\n        x = self.q_sample(x_start = x_start, t = t, noise = noise)\n\n        # predict and take gradient step\n\n        model_out = self.model(x, t)\n\n        loss = self.loss_fn(model_out, noise, reduction = 'none')\n        loss = reduce(loss, 'b ... -> b (...)', 'mean')\n\n        loss = loss * extract(self.loss_weight, t, loss.shape)\n        return loss.mean()\n\n    def forward(self, img, *args, **kwargs):\n        b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size\n        assert h == img_size and w == img_size, f'height and width of image must be {img_size}'\n        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()\n\n        img = self.normalize(img)\n        return self.p_losses(img, t, *args, **kwargs)\n<\/code><\/pre>\n\n\n\n
Define Trainer: define the updating process<\/p>\n\n\n\n
class Trainer(object):\n    def __init__(\n        self,\n        diffusion_model,\n        folder,\n        *,\n        train_batch_size = 16,\n        gradient_accumulate_every = 1,\n        train_lr = 1e-4,\n        train_num_steps = 100000,\n        ema_update_every = 10,\n        ema_decay = 0.995,\n        adam_betas = (0.9, 0.99),\n        save_and_sample_every = 1000,\n        num_samples = 25,\n        results_folder = '.\/results',\n        split_batches = True,\n        inception_block_idx = 2048\n    ):\n        super().__init__()\n\n        # accelerator\n\n        self.accelerator = Accelerator(\n            split_batches = split_batches,\n            mixed_precision = 'no'\n        )\n\n\n        # model\n\n        self.model = diffusion_model\n        self.channels = diffusion_model.channels\n\n        # sampling and training hyperparameters\n\n        assert has_int_squareroot(num_samples), 'number of samples must have an integer square root'\n        self.num_samples = num_samples\n        self.save_and_sample_every = save_and_sample_every\n\n        self.batch_size = train_batch_size\n        self.gradient_accumulate_every = gradient_accumulate_every\n\n        self.train_num_steps = train_num_steps\n        self.image_size = diffusion_model.image_size\n\n        # dataset and dataloader\n\n        self.ds = Dataset(folder, self.image_size)\n        dl = DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count())\n\n        dl = self.accelerator.prepare(dl)\n        self.dl = cycle(dl)\n\n        # optimizer\n\n        self.opt = Adam(diffusion_model.parameters(), lr = train_lr, betas = adam_betas)\n\n        # for logging results in a folder periodically\n\n        if self.accelerator.is_main_process:\n            self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)\n            self.ema.to(self.device)\n\n        self.results_folder = Path(results_folder)\n        self.results_folder.mkdir(exist_ok = True)\n\n        # step counter state\n\n        self.step = 0\n\n        # prepare model, dataloader, optimizer with accelerator\n\n        self.model, self.opt = self.accelerator.prepare(self.model, self.opt)\n\n    @property\n    def device(self):\n        return self.accelerator.device\n\n    def save(self, milestone):\n        if not self.accelerator.is_local_main_process:\n            return\n\n        data = {\n            'step': self.step,\n            'model': self.accelerator.get_state_dict(self.model),\n            'opt': self.opt.state_dict(),\n            'ema': self.ema.state_dict(),\n            'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None,\n        }\n\n        torch.save(data, str(self.results_folder \/ f'model-{milestone}.pt'))\n\n    def load(self, ckpt):\n        accelerator = self.accelerator\n        device = accelerator.device\n\n        data = torch.load(ckpt, map_location=device)\n\n        model = self.accelerator.unwrap_model(self.model)\n        model.load_state_dict(data['model'])\n\n        self.step = data['step']\n        self.opt.load_state_dict(data['opt'])\n        if self.accelerator.is_main_process:\n            self.ema.load_state_dict(data[\"ema\"])\n\n\n        if exists(self.accelerator.scaler) and exists(data['scaler']):\n            self.accelerator.scaler.load_state_dict(data['scaler'])\n\n\n    def train(self):\n        accelerator = self.accelerator\n        device = accelerator.device\n\n        with tqdm(initial = self.step, total = self.train_num_steps, disable = not accelerator.is_main_process) as pbar:\n\n            while self.step < self.train_num_steps:\n\n                total_loss = 0.\n\n                for _ in range(self.gradient_accumulate_every):\n                    data = next(self.dl).to(device)\n\n                    with self.accelerator.autocast():\n                        loss = self.model(data)\n                        loss = loss \/ self.gradient_accumulate_every\n                        total_loss += loss.item()\n\n                    self.accelerator.backward(loss)\n\n                accelerator.clip_grad_norm_(self.model.parameters(), 1.0)\n                pbar.set_description(f'loss: {total_loss:.4f}')\n\n                accelerator.wait_for_everyone()\n\n                self.opt.step()\n                self.opt.zero_grad()\n\n                accelerator.wait_for_everyone()\n\n                self.step += 1\n                if accelerator.is_main_process:\n                    self.ema.update()\n\n                    if self.step != 0 and self.step % self.save_and_sample_every == 0:\n                        self.ema.ema_model.eval()\n\n                        with torch.no_grad():\n                            milestone = self.step \/\/ self.save_and_sample_every\n                            batches = num_to_groups(self.num_samples, self.batch_size)\n                            all_images_list = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches))\n\n                        all_images = torch.cat(all_images_list, dim = 0)\n\n                        utils.save_image(all_images, str(self.results_folder \/ f'sample-{milestone}.png'), nrow = int(math.sqrt(self.num_samples)))\n\n                        self.save(milestone)\n\n                pbar.update(1)\n\n        accelerator.print('training complete')\n\n    def inference(self, num=1000, n_iter=5, output_path='.\/submission'):\n        if not os.path.exists(output_path):\n            os.mkdir(output_path)\n        with torch.no_grad():\n            for i in range(n_iter):\n                batches = num_to_groups(num \/\/ n_iter, 200)\n                all_images = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches))[0]\n                for j in range(all_images.size(0)):\n                    torchvision.utils.save_image(all_images[j], f'{output_path}\/{i * 200 + j + 1}.jpg')\n<\/code><\/pre>\n\n\n\n
Training Hyper-parameters<\/p>\n\n\n\n
path = '\/content\/faces'\nIMG_SIZE = 64             # Size of images, do not change this if you do not know why you need to change\nbatch_size = 16\ntrain_num_steps = 10000        # total training steps\nlr = 1e-3\ngrad_steps = 1            # gradient accumulation steps, the equivalent batch size for updating equals to batch_size * grad_steps = 16 * 1\nema_decay = 0.995           # exponential moving average decay\n\nchannels = 16             # Numbers of channels of the first layer of CNN\ndim_mults = (1, 2, 4)        # The model size will be (channels, 2 * channels, 4 * channels, 4 * channels, 2 * channels, channels)\n\ntimesteps = 100            # Number of steps (adding noise)\nbeta_schedule = 'linear'\n\nmodel = Unet(\n    dim = channels,\n    dim_mults = dim_mults\n)\n\ndiffusion = GaussianDiffusion(\n    model,\n    image_size = IMG_SIZE,\n    timesteps = timesteps,\n    beta_schedule = beta_schedule\n)\n\ntrainer = Trainer(\n    diffusion,\n    path,\n    train_batch_size = batch_size,\n    train_lr = lr,\n    train_num_steps = train_num_steps,\n    gradient_accumulate_every = grad_steps,\n    ema_decay = ema_decay,\n    save_and_sample_every = 1000\n)\n\ntrainer.train()<\/code><\/pre>\n\n\n\n
\/usr\/local\/lib\/python3.10\/dist-packages\/accelerate\/accelerator.py:436: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['split_batches']). Please pass an `accelerate.DataLoaderConfiguration` instead: \ndataloader_config = DataLoaderConfiguration(split_batches=True)\n  warnings.warn(\n<\/pre>\n\n\n\n
loss:\u20070.0785:\u2007\u200766%<\/p>\n\n\n\n
\u20076569\/10000\u2007[6:44:31<3:18:41,\u2007\u20073.47s\/it]<\/p>\n\n\n\n
\/usr\/lib\/python3.10\/multiprocessing\/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n  self.pid = os.fork()\n<\/pre>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[02:11<00:00,\u2007\u20071.26s\/it]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[01:01<00:00,\u2007\u20071.73it\/s]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[02:22<00:00,\u2007\u20071.36s\/it]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[01:01<00:00,\u2007\u20071.74it\/s]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[02:18<00:00,\u2007\u20071.33s\/it]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[01:00<00:00,\u2007\u20071.74it\/s]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[02:16<00:00,\u2007\u20071.34s\/it]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[01:00<00:00,\u2007\u20071.72it\/s]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[02:12<00:00,\u2007\u20071.32s\/it]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[01:01<00:00,\u2007\u20071.65it\/s]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[02:21<00:00,\u2007\u20071.50s\/it]<\/p>\n\n\n\n
sampling\u2007loop\u2007time\u2007step:\u2007100%<\/p>\n\n\n\n
\u2007100\/100\u2007[01:03<00:00,\u2007\u20071.46it\/s]<\/p>\n\n\n\n
\"\"<\/div><\/figure>\n\n\n\n
Inference<\/p>\n\n\n\n
ckpt = '\/content\/drive\/MyDrive\/ML 2023 Spring\/model-55.pt'\ntrainer.load(ckpt)\ntrainer.inference()<\/code><\/pre>\n\n\n\n
%cd .\/submission\n!tar -zcf ..\/submission.tgz *.jpg\n%cd ..<\/code><\/pre>\n\n\n\n
\n","protected":false},"excerpt":{"rendered":"
\u3010HW6\u3011Diffusion Model0.1\u674e\u5b8f\u6bc52021\/2022\u6625\u673a\u5668\u5b66\u4e60\u8bfe\u7a0b\u7b14\u8bb0EP18(P68-P7 […]<\/p>\n","protected":false},"author":1,"featured_media":0,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[6],"tags":[15,3,7,9,8],"class_list":["post-664","post","type-post","status-publish","format-standard","hentry","category-lhyjqxxbj","tag-homework","tag-xxbj","tag-jjxx","tag-lhy","tag-deeplearning"],"_links":{"self":[{"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/posts\/664","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/comments?post=664"}],"version-history":[{"count":2,"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/posts\/664\/revisions"}],"predecessor-version":[{"id":1873,"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/posts\/664\/revisions\/1873"}],"wp:attachment":[{"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/media?parent=664"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/categories?post=664"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/tobykskgd.life\/index.php\/wp-json\/wp\/v2\/tags?post=664"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}