Function test_fn

diffusion_engine.py:352–505

(
    test_data_path,
    weight_path,
    schedule_type="cosine",
    batch_size=320,
    n_steps=1500,
    show=False,
    device="cuda:0",
    full_res=False,
    dataset_name="gf2",
    division=1023,
)

@torch.no_grad()
def test_fn(
    test_data_path,
    weight_path,
    schedule_type="cosine",
    batch_size=320,
    n_steps=1500,
    show=False,
    device="cuda:0",
    full_res=False,
    dataset_name="gf2",
    division=1023,
):
    # lazy import
    from models.sr3_dwt import UNetSR3 as Unet
    from diffusion.diffusion_ddpm_pan import GaussianDiffusion

    torch.cuda.set_device(device)

    # load model
    if dataset_name in ['harvard', 'cave']:
        image_n_channel = 31
        image_size = 512 if dataset_name == 'cave' else 1000
        pan_channel = 3
        rgb_channels = [39, 19, 9]
    elif dataset_name in ['wv3', 'gf2', 'qb']:
        image_size = 512 if full_res else 256
        image_n_channel = 8 if dataset_name == 'wv3' else 4
        pan_channel = 1
        rgb_channels = [4, 2, 0] if dataset_name == 'wv3' else [2, 1, 0]
    denoise_fn = Unet(
        in_channel=image_n_channel,
        out_channel=image_n_channel,
        lms_channel=image_n_channel,
        pan_channel=pan_channel,#1,
        inner_channel=32,  # 32,
        norm_groups=1,
        channel_mults=(1, 2, 2, 4),  # (64, 32, 16, 8)
        attn_res=(8,),
        dropout=0.2,
        image_size=64,
        self_condition=True,
    ).to(device)
    denoise_fn = model_load(weight_path, denoise_fn, device=device)

    denoise_fn.eval()
    print(f"load weight {weight_path}")
    diffusion = GaussianDiffusion(
        denoise_fn,
        image_size=image_size,
        channels=image_n_channel,
        pred_mode="x_start",
        loss_type="l1",
        device=device,
        clamp_range=(0, 1),
    )
    diffusion.set_new_noise_schedule(
        betas=make_beta_schedule(schedule="cosine", n_timestep=n_steps, cosine_s=8e-3)
    )