Python make_beta_schedule Examples

Example #1

File: train.py Project: its-dron/denoising-diffusion-pytorch

def main(conf):
    device = "cuda:0" if torch.cuda.is_available() else 'cpu'
    beta_schedule = "linear"
    beta_start = 1e-4
    beta_end = 2e-2
    n_timestep = 1000

    conf.distributed = dist.get_world_size() > 1

    transform = transforms.Compose(
        [
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )

    train_set = MultiResolutionDataset(
        conf.dataset.path, transform, conf.dataset.resolution
    )
    train_sampler = dist.data_sampler(
        train_set, shuffle=True, distributed=conf.distributed
    )
    train_loader = conf.training.dataloader.make(train_set, sampler=train_sampler)

    model = UNet(
        conf.model.in_channel,
        conf.model.channel,
        channel_multiplier=conf.model.channel_multiplier,
        n_res_blocks=conf.model.n_res_blocks,
        attn_strides=conf.model.attn_strides,
        dropout=conf.model.dropout,
        fold=conf.model.fold,
    )
    model = model.to(device)
    ema = UNet(
        conf.model.in_channel,
        conf.model.channel,
        channel_multiplier=conf.model.channel_multiplier,
        n_res_blocks=conf.model.n_res_blocks,
        attn_strides=conf.model.attn_strides,
        dropout=conf.model.dropout,
        fold=conf.model.fold,
    )
    ema = ema.to(device)

    if conf.distributed:
        model = nn.parallel.DistributedDataParallel(
            model,
            device_ids=[dist.get_local_rank()],
            output_device=dist.get_local_rank(),
        )

    optimizer = conf.training.optimizer.make(model.parameters())
    scheduler = conf.training.scheduler.make(optimizer)

    betas = make_beta_schedule(beta_schedule, beta_start, beta_end, n_timestep)
    diffusion = GaussianDiffusion(betas).to(device)

    train(conf, train_loader, model, ema, diffusion, optimizer, scheduler, device)

Example #2

File: generate_sample.py Project: its-dron/denoising-diffusion-pytorch

def generate_sample(model):
    betas = make_beta_schedule('linear', 1e-4, 2e-2, 1000)
    diffusion = GaussianDiffusion(betas).to('cuda')

    imgs = p_sample_loop(diffusion, model, [16, 3, 128, 128], 'cuda', capture_every=10)
    imgs = imgs[1:]

    id = 0
    grid = make_grid(torch.cat([i[id:id + 1] for i in imgs[:-1:4]], 0), nrow=5, normalize=True, range=(-1, 1))

    return grid.detach().mul(255).cpu().type(torch.uint8).permute(1, 2, 0).numpy()

Example #3

File: test_op.py Project: HighCWu/denoising-diffusion-paddle

import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from model import UNet
from diffusion import make_beta_schedule, GaussianDiffusion
from config import config

conf = config.diffusion
betas = make_beta_schedule(**conf.diffusion.beta_schedule)
diffusion = GaussianDiffusion(betas)
model = UNet(**conf.model)
img = paddle.randn([
    conf.training.dataloader.batch_size,
    conf.model.in_channel,
    conf.dataset.resolution,
    conf.dataset.resolution
])
time = paddle.randint(
    0, conf.diffusion.beta_schedule.n_timestep, (img.shape[0],)
)
loss = diffusion.p_loss(model, img, time)
print(loss.numpy())


conf = config.improved
betas = make_beta_schedule(**conf.diffusion.beta_schedule)
diffusion = GaussianDiffusion(betas)
model = UNet(**conf.model)
img = paddle.randn([
    conf.training.dataloader.batch_size,