分享来自"我需要多少次迭代?"的答案

    dim = 64,
    dim_mults = (1, 2, 4, 8)
).cuda()

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,           # 步数
    sampling_timesteps = 250,   # 采样步数
    loss_type = 'l1'            # L1 或 L2
).cuda()

trainer = Trainer(
    diffusion,
    '/mnt/prj/seungyoun/dataset/flowers',
    train_batch_size = 128,
    train_lr = 1e-4,
    train_num_steps = 700000,         # 总训练步数
    gradient_accumulate_every = 2,    # 累积梯度步数
    ema_decay = 0.995,                # 指数移动平均衰减
    amp = False                       # 开启混合精度
)

trainer.train()