DynamicDPS/image_train.py at master · edshkim98/DynamicDPS · 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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
"""
Train a diffusion model on images.
"""

import argparse

from guided_diffusion import logger
from guided_diffusion.image_datasets import load_data, IQTDataset
from guided_diffusion.resample import create_named_schedule_sampler
from guided_diffusion.script_util import (
    model_and_diffusion_defaults,
    create_model_and_diffusion,
    args_to_dict,
    add_dict_to_argparser,
)
from guided_diffusion.train_util import TrainLoop
from torch.utils.data import Dataset, DataLoader
import yaml
import glob
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_data_custom(data_loader):
    while True:
        yield from data_loader

def main():

    with open('/cluster/project0/IQT_Nigeria/skim/diffusion_inverse/guided-diffusion/configs.yaml') as file:
        configs = yaml.load(file, Loader=yaml.FullLoader)
    args = create_argparser().parse_args()

    # dist_util.setup_dist()
    logger.configure(dir="./logs_large_zero2two_HCPMoreSlice2025")

    logger.log("creating model and diffusion...")
    model, diffusion = create_model_and_diffusion(configs = configs,
        **args_to_dict(args, model_and_diffusion_defaults().keys())
    )
    if configs['fine_tune']:
            model.load_state_dict(
        torch.load("./logs_large_zero2two/model120000.pt", map_location="cpu")
    )

    model = model.to(device)
    schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion)

    logger.log("creating data loader...")
    # data = load_data(
    #     data_dir=args.data_dir,
    #     batch_size=args.batch_size,
    #     image_size=args.image_size,
    #     class_cond=args.class_cond,
    # )

    files = glob.glob(args.data_dir[0] + '/T1w/T1w_acpc_dc_restore_brain.nii.gz')#T1w_acpc_dc_restore_brain_sim036T_4x_groundtruth.nii.gz')
    print("File being used: ", args.data_dir[0])
    if len(args.data_dir[0]) > 1:
        for i in range(len(args.data_dir)):
            if 'skim' in args.data_dir[i]:
                files_new = glob.glob(args.data_dir[i] + '/gt.npy')
            else:
                files_new = glob.glob(args.data_dir[i] + '/T1w/T1w_acpc_dc_restore_brain.nii.gz')
    dataset = IQTDataset(files, configs = configs)
    print(f"Files: {len(files)} Dataset size: {len(dataset)}")
    data = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=1, drop_last=False)
    data = load_data_custom(data)

    logger.log("training...")
    TrainLoop(
        model=model,
        diffusion=diffusion,
        data=data,
        batch_size=args.batch_size,
        microbatch=args.microbatch,
        lr=args.lr,
        ema_rate=args.ema_rate,
        log_interval=args.log_interval,
        save_interval=args.save_interval,
        resume_checkpoint=args.resume_checkpoint,
        use_fp16=args.use_fp16,
        fp16_scale_growth=args.fp16_scale_growth,
        schedule_sampler=schedule_sampler,
        weight_decay=args.weight_decay,
        lr_anneal_steps=args.lr_anneal_steps,
    ).run_loop()


def create_argparser():
    defaults = dict(
        data_dir=["/cluster/project0/IQT_Nigeria/HCP_t1t2_ALL/sim/[!9]*/"],#"/cluster/project0/IQT_Nigeria/skim/HCP_Harry_x4/train/",
        schedule_sampler="uniform",
        lr=1e-4,
        weight_decay=0.0,
        lr_anneal_steps=0,
        batch_size=8,
        microbatch=-1,  # -1 disables microbatches
        ema_rate="0.9999",  # comma-separated list of EMA values
        log_interval=10,
        save_interval=10000,
        resume_checkpoint="",
        use_fp16=False,
        fp16_scale_growth=1e-3,
    )
    defaults.update(model_and_diffusion_defaults())
    parser = argparse.ArgumentParser()
    add_dict_to_argparser(parser, defaults)
    return parser


if __name__ == "__main__":
    main()