Fork me on GitHub

DDPM代码详细解读(1):数据集准备、超参数设置、loss设计、关键参数计算

字数统计: 1.4k字   |   阅读时长≈ 7分

DDPM代码详细解读(1):数据集准备、超参数设置、loss设计、关键参数计算

之前照猫画虎复现了DDPM，但是其中许多代码很不好理解，遂写此博客进行学习记录。

超参数设置

超参数设置使用absl中flags进行管理，num_res_blocks是U-net中每个level中resnet层数，attn是attention block，这个是后面我们加入condition的途径，非常重要。

beta_1和beta_2对用于\(\beta _1\)和\(\beta _T\)，实际的\(\beta _t\)是在\(\beta _1\)、\(\beta _T\)中线性采样得到的，DDPM原文中研究了是否固定\(\beta _t\)对实验结果的影响，后面很多论文也做了对比实验探索是否线性对实验效果的影响。(个人感觉这个参数是否线性对实验的影响不大，只需要满足足够小的假设即可)

T是采样的步长，这个对采样质量和生成时间影响非常大。(在实验中，T越大，采样时间越长，4070显卡采样一个batch的数据设置需要近20小时。但是T越大并不是质量越高，呈二次函数关系。)

image_size根据数据集实际情况设置，这是影响生成时间的重要因素，size和时间呈指数倍爆炸增长。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
FLAGS = flags.FLAGS
flags.DEFINE_bool('train', False, help='train from scratch')
flags.DEFINE_bool('eval', False, help='load model.pt and evaluate FID and IS')
# UNet
flags.DEFINE_integer('ch', 128, help='base channel of UNet')
flags.DEFINE_multi_integer('ch_mult', [1, 2, 2, 2], help='channel multiplier')
flags.DEFINE_multi_integer('attn', [1], help='add attention to these levels')
flags.DEFINE_integer('num_res_blocks', 2, help='# resblock in each level')
flags.DEFINE_float('dropout', 0.1, help='dropout rate of resblock')
# Gaussian Diffusion
flags.DEFINE_float('beta_1', 1e-4, help='start beta value')
flags.DEFINE_float('beta_T', 0.02, help='end beta value')
flags.DEFINE_integer('T', 1000, help='total diffusion steps')
flags.DEFINE_enum('mean_type', 'epsilon', ['xprev', 'xstart', 'epsilon'], help='predict variable')
flags.DEFINE_enum('var_type', 'fixedlarge', ['fixedlarge', 'fixedsmall'], help='variance type')
# Training
flags.DEFINE_float('lr', 2e-4, help='target learning rate')
flags.DEFINE_float('grad_clip', 1., help="gradient norm clipping")
flags.DEFINE_integer('total_steps', 800000, help='total training steps')
flags.DEFINE_integer('img_size', 32, help='image size')
flags.DEFINE_integer('warmup', 5000, help='learning rate warmup')
flags.DEFINE_integer('batch_size', 128, help='batch size')
flags.DEFINE_integer('num_workers', 4, help='workers of Dataloader')
flags.DEFINE_float('ema_decay', 0.9999, help="ema decay rate")
flags.DEFINE_bool('parallel', False, help='multi gpu training')
# Logging & Sampling
flags.DEFINE_string('logdir', './logs/DDPM_CIFAR10_EPS', help='log directory')
flags.DEFINE_integer('sample_size', 64, "sampling size of images")
flags.DEFINE_integer('sample_step', 1000, help='frequency of sampling')
# Evaluation
flags.DEFINE_integer('save_step', 5000, help='frequency of saving checkpoints, 0 to disable during training')
flags.DEFINE_integer('eval_step', 0, help='frequency of evaluating model, 0 to disable during training')
flags.DEFINE_integer('num_images', 50000, help='the number of generated images for evaluation')
flags.DEFINE_bool('fid_use_torch', False, help='calculate IS and FID on gpu')
flags.DEFINE_string('fid_cache', './stats/cifar10.train.npz', help='FID cache')
 
device = torch.device('cuda:0')

训练CIFAR10数据集的配置信息

由于不同的数据集U-net channel、T、image_size等关键参数是不一样的，因此针对不同的数据集采用不同的txt文件进行管理。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
--T=1000
--attn=1
--batch_size=128
--beta_1=0.0001
--beta_T=0.02
--ch=128
--ch_mult=1
--ch_mult=2
--ch_mult=2
--ch_mult=2
--dropout=0.1
--ema_decay=0.9999
--noeval
--eval_step=0
--fid_cache=./stats/cifar10.train.npz
--nofid_use_torch
--grad_clip=1.0
--img_size=32
--logdir=./logs/DDPM_CIFAR10_EPS
--lr=0.0002
--mean_type=epsilon
--num_images=50000
--num_res_blocks=2
--num_workers=4
--noparallel
--sample_size=64
--sample_step=1000
--save_step=5000
--total_steps=800000
--train
--var_type=fixedlarge
--warmup=5000

加载数据集

以加载cifar10数据集为例：

1
2
3
4
5
6
7
8
9
10
11
12
# dataset
dataset = CIFAR10(
    root='./data', train=True, download=True,
    transform=transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]))
dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=FLAGS.batch_size, shuffle=True,
    num_workers=FLAGS.num_workers, drop_last=True)
datalooper = infiniteloop(dataloader)

每个loop使用next方法即可：

1
x_0 = next(datalooper).to(device)

loss计算

原文的loss计算公式： \[ L_{simple}\left( \theta \right) :=\mathbb{E}_{t,x_0,\epsilon}\left[ ||\epsilon -\epsilon _{\theta}\left( \sqrt{\bar{\alpha}}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon ,t \right) ||\begin{array}{c} 2\\ \\\ \end{array} \right] \] 计算的是纯噪声和\(\epsilon _{\theta}\left( \sqrt{\bar{\alpha}}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon ,t \right)\)之间的均方差：

1
loss = F.mse_loss(self.model(x_t, t), noise, reduction='none')

其中noise的size和input data一样的：

1
noise = torch.randn_like(x_0)

计算\(x_t\)

具体的计算公式如下： \[ x_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon \] 其中时刻信息 t 是通过\(\bar{\alpha}_t\)表现的，代码如下：

1
2
3
4
5
6
7
8
9
10
def extract(v, t, x_shape):
    """
    Extract some coefficients at specified timesteps, then reshape to
    [batch_size, 1, 1, 1, 1, ...] for broadcasting purposes.
    """
    out = torch.gather(v, index=t, dim=0).float()
    return out.view([t.shape[0]] + [1] * (len(x_shape) - 1))       
 x_t = (
            extract(self.sqrt_alphas_bar, t, x_0.shape) * x_0 +
            extract(self.sqrt_one_minus_alphas_bar, t, x_0.shape) * noise)

其中extract函数的作用是选取特定下标 t 的信息并转换成特定维度。

计算\(\bar{\alpha}_t\)

1、根据\(\beta _1\)和\(\beta _T\)计算所有的\(\beta _t\)

DDPM原始的论文设置的\(\beta _t\)是线性增长，后面不少文章设置了指数增长等其他方式，只要满足足够小假设即可。

1
2
self.register_buffer(
    'betas', torch.linspace(beta_1, beta_T, T).double())

2、计算\(\alpha_t\) \[ \alpha _t=1-\beta _t \]

1
alphas = 1. - self.betas

3、累乘得到\(\bar{\alpha}_t\)

1
alphas_bar = torch.cumprod(alphas, dim=0)

最后将这些一同写入buffer即可：

1
2
3
4
self.register_buffer(
    'sqrt_alphas_bar', torch.sqrt(alphas_bar))
self.register_buffer(
    'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar)

计算\(\frac{1-\alpha _t}{\sqrt{1-\bar{\alpha}_t}}\)

image-20231113173032381

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
    self.register_buffer(
        'posterior_log_var_clipped',
        torch.log(
            torch.cat([self.posterior_var[1:2], self.posterior_var[1:]])))
    self.register_buffer(
        'posterior_mean_coef1',
        torch.sqrt(alphas_bar_prev) * self.betas / (1. - alphas_bar))
    self.register_buffer(
        'posterior_mean_coef2',
        torch.sqrt(alphas) * (1. - alphas_bar_prev) / (1. - alphas_bar))
def predict_xstart_from_xprev(self, x_t, t, xprev):
    assert x_t.shape == xprev.shape
    return (  # (xprev - coef2*x_t) / coef1
        extract(
            1. / self.posterior_mean_coef1, t, x_t.shape) * xprev -
        extract(
            self.posterior_mean_coef2 / self.posterior_mean_coef1, t,
            x_t.shape) * x_t
    )

扫一扫，分享到微信

tag:

    缺失模块。
    1、请确保node版本大于6.2
    2、在博客根目录（注意不是yilia-plus根目录）执行以下命令：
    npm i hexo-generator-json-content --save

    3、在根目录_config.yml里添加配置： 

      jsonContent:
    meta: false
    pages: false
    posts:
      title: true
      date: true
      path: true
      text: false
      raw: false
      content: false
      slug: false
      updated: false
      comments: false
      link: false
      permalink: false
      excerpt: false
      categories: false
      tags: true

计算机小学徒，慢慢成长