train_diff.py

import os
import numpy as np
import torchvision
import torch
import time
import argparse
import torchvision.transforms as transforms
from torch.optim import Adam
from models.denoiseModel.UNet import *
from models.diffusionModel.Diffusion import *
from utils.trainer import *
from matplotlib import pyplot as plt
from dataloaders.datasets import *
from visdom import Visdom
from torch.utils.data import DataLoader


def get_parser():
    parser = argparse.ArgumentParser()

    parser.add_argument("--batch_size", type=int, default=1, help='Training batch size')
    parser.add_argument("--channels", type=int, default=1, help='Number of channels')
    parser.add_argument("--img_size", type=int, default=240, help='Image size (will be resized to this size)')
    parser.add_argument("--dax", type=str, default="BraTS18", choices=['Abdominal', 'BraTS18', 'Fundus'],
                        help='Dataset name')
    parser.add_argument("--data_dir", type=str,
                        default="./datasets/BraTS_2018", help='Dataset directory')
    parser.add_argument("--data_transform", help='Whether to preprocess the data', action='store_true')

    parser.add_argument("--denoise_mode", type=str, default="UNetFree", help='Denoising model type')
    parser.add_argument("--dim_mults", type=tuple, default=(1, 2, 4, 8),
                        help='Channel multipliers for each model stage')
    parser.add_argument("--dim_bn", type=int, default=32, help='Base channel dimension for each model stage')
    parser.add_argument("--betas_schedule_name", type=str, default="linear", help='Noise schedule')
    parser.add_argument("--time_steps", type=int, default=1000, help='Number of diffusion timesteps')
    parser.add_argument("--beta_start", type=float, default=0.0001, help='Starting beta (noise variance)')
    parser.add_argument("--beta_end", type=float, default=0.02, help='Ending beta (noise variance)')
    parser.add_argument("--classifier_mode", type=str, default="free", help='Conditional guidance mode')
    parser.add_argument("--root_path", type=str, default="./saved_train_models/diff_model", help='Save path')
    parser.add_argument("--noise_mu", type=float, default=1.0, help='Noise scale / noise base factor')

    parser.add_argument("--epochs", type=int, default=100, help='Number of training epochs')
    parser.add_argument("--lr_diff", type=float, default=1e-4, help='Denoising model learning rate')
    parser.add_argument("--save_epochs", type=int, default=None, help='Model saving interval (in epochs)')
    parser.add_argument("--train_log", help='Whether to train on log-transformed data (enable to train in log space)',
                        action='store_true')
    parser.add_argument("--all_vis_loss", help='Whether to send loss from every iteration to Visdom',
                        action='store_true')
    parser.add_argument("--source_domain", type=str, default="t2", help='Source domain')
    parser.add_argument("--loss_type", type=str, default="huber", help='Loss type')
    parser.add_argument("--ifft_flg", help='Whether to apply inverse transform to cond', action='store_true')

    parser.add_argument("--visdom_flg", help='Enable visualization via Visdom', action='store_true')
    parser.add_argument("--visdom_port", type=int, default=8097, help='Visdom port')
    parser.add_argument("--tqdm_flg", help='Enable tqdm progress bar', action='store_true')

    parser.add_argument("--allsave_path", type=str, default="./log/allsave.txt",
                        help='Log file for args and corresponding paths')

    parser.add_argument('--gpus', type=str, default='0', help='GPU IDs to use (comma-separated, e.g., 0,1)')

    args = parser.parse_args()
    show_args(args)
    return args


def show_args(args):
    for (key, value) in args.__dict__.items():
        print(key, value)


def get_cmd(args):
    strcmd = 'python train_diff.py'
    for (key, value) in args.__dict__.items():
        if value is False:
            continue
        if value is True:
            strcmd += f' --{key}'
        else:
            strcmd += f' --{key} {value}'
    return strcmd


def args_on_visdom(vis, args):
    vis.text(f'args', win='args_text')
    for (key, value) in args.__dict__.items():
        vis.text(f'{key}: {value}', win='args_text', append=True)
    vis.text(f'==============================', win='args_text', append=True)
    vis.text(f'start by:', win='args_text', append=True)
    vis.text(f'{get_cmd(args)}', win='args_text', append=True)


def write_args(id_path, args):
    directory = os.path.dirname(args.allsave_path)
    os.makedirs(directory, exist_ok=True)
    with open(args.allsave_path, 'a+') as fi:
        fi.write(f'{id_path}\n')
        fi.write('------------------------------\n')
        for (key, value) in args.__dict__.items():
            fi.write(f'{key}: {value}\n')
        fi.write('------------------------------\n')
        fi.write(get_cmd(args))
        fi.write('\n')
        fi.write('------------------------------\n')
        fi.write(args.root_path + '/' + str(id_path))
        fi.write('\n')
        fi.write('==============================\n')


def run(args):
    # cuda
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if not os.path.exists(args.root_path):
        os.makedirs(args.root_path)
    save_path = None
    id_path = 0
    for id_path in range(0, 10000):
        save_path = args.root_path + '/' + str(id_path)
        if not os.path.exists(args.root_path + '/' + str(id_path)):
            os.makedirs(args.root_path + '/' + str(id_path))
            write_args(id_path, args)
            break

    if args.visdom_flg:
        vis = Visdom(env=f'diffusion_train_{id_path}', port=args.visdom_port)
        if not vis.check_connection():
            raise ConnectionError("Visdom server is not running")
        args_on_visdom(vis, args)

    def worker_init_fn(worker_id):
        random.seed(args.seed + worker_id)

    if not os.path.exists(args.data_dir):
        os.makedirs(args.data_dir)

    tran_list = None
    tran_train = None
    if args.data_transform:
        tran_list = [transforms.Resize((args.img_size, args.img_size)), transforms.ToTensor()]
        tran_train = transforms.Compose(tran_list)

    if args.dax == 'BraTS18':
        dataset = BraTS18Dataset(args.data_dir, dana=get_dom(args=args))
    elif args.dax == 'Fundus':
        dataset = FundusDataset(args.data_dir, dana=get_dom(args=args), img_size=args.img_size)
    else:
        raise NotImplementedError
    data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, worker_init_fn=worker_init_fn)

    # # denoise model
    denoise_model = None
    if args.denoise_mode == 'UNetFree':
        denoise_model = UNetFree(in_dim=args.dim_bn, dim_mults=args.dim_mults, channels=args.channels,
                                 cond_channel=args.channels).to(device)
    else:
        raise NotImplementedError
    # # diffusion tool
    diffusion = Diffusion(shape=(args.channels, args.img_size, args.img_size), denoise_model=denoise_model,
                          betas_schedule_name=args.betas_schedule_name, time_steps=args.time_steps,
                          beta_start=args.beta_start, \
                          beta_end=args.beta_end, classifier_mode=args.classifier_mode, rand_mu=args.noise_mu)
    # # optimizer
    optimizer = Adam(diffusion.parameters(), lr=args.lr_diff)

    trainer = FFTCondDiffusionTrainer(epoches=args.epochs,
                                      train_loader=data_loader,
                                      optimizer=optimizer,
                                      device=device,
                                      time_steps=args.time_steps,
                                      IFvisdom=args.visdom_flg,
                                      visdom_port=args.visdom_port,
                                      Visdom_env=f'FFTCDiff_{id_path}',
                                      save_epochs=args.save_epochs,
                                      tqdm_flg=args.tqdm_flg,
                                      log_flg=args.train_log,
                                      all_vis_loss=args.all_vis_loss,
                                      loss_type=args.loss_type,
                                      ifft_flg=args.ifft_flg,
                                      )

    diffusion = trainer(diffusion, model_save_path=save_path)


if __name__ == '__main__':
    run(get_parser())