PaD-TS/diffmodel_init.py at main · wmd3i/PaD-TS · GitHub

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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70


import argparse
import inspect
import numpy as np
import torch as th

import gaussian_diffusion_loss as gd
from respace import space_timesteps, SpacedDiffusion


def diffusion_defaults():
    """
    Defaults for image training.
    """
    return dict(
        learn_sigma=False,
        sigma_small=False,
        diffusion_steps=1000,
        noise_schedule="linear",
        timestep_respacing="",
        use_kl=False,
        predict_xstart=False,
        rescale_timesteps=True,
        rescale_learned_sigmas=True,
    )


def create_gaussian_diffusion(
    *,
    diffusion_steps=1000,
    learn_sigma=False,
    sigma_small=False,
    noise_schedule="linear",
    loss='MSE',
    predict_xstart=False,
    rescale_timesteps=False,
    timestep_respacing="",
):
    betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
    assert loss in ['MSE', 'MSE_MMD', 'MSE_CORR']

    if loss == 'MSE_MMD':
        loss_type = gd.LossType.MSE_MMD
    elif loss == 'MSE_CORR':
        loss_type = gd.LossType.MSE_CORR
    else:
        loss_type = gd.LossType.MSE
    if not timestep_respacing:
        timestep_respacing = [diffusion_steps]
    return SpacedDiffusion(
        use_timesteps=space_timesteps(diffusion_steps, timestep_respacing),
        betas=betas,
        model_mean_type=(
            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
        ),
        model_var_type=(
            (
                gd.ModelVarType.FIXED_LARGE
                if not sigma_small
                else gd.ModelVarType.FIXED_SMALL
            )
            if not learn_sigma
            else gd.ModelVarType.LEARNED_RANGE
        ),
        loss_type=loss_type,
        rescale_timesteps=rescale_timesteps,
    )