train and valid function Args: train_dataset_path (str): _description_ valid_dataset_path (str): _description_ batch_size (int, optional): _description_. Defaults to 240. n_steps (int, optional): _description_. Defaults to 1500. epochs (int, optional): _d
(
# dataset
train_dataset_path,
valid_dataset_path,
dataset_name=None,
# image settings
image_n_channel=8,
image_size=64,
# diffusion settings
schedule_type="cosine",
n_steps=3_000,
max_iterations=400_000,
# device setting
device="cuda:0",
# optimizer settings
batch_size=128,
lr_d=1e-4,
show_recon=False,
# pretrain settings
pretrain_weight=None,
pretrain_iterations=None,
*,
# just for debugging
constrain_channel=None,
)
| 50 | |
| 51 | |
| 52 | def engine_google( |
| 53 | # dataset |
| 54 | train_dataset_path, |
| 55 | valid_dataset_path, |
| 56 | dataset_name=None, |
| 57 | # image settings |
| 58 | image_n_channel=8, |
| 59 | image_size=64, |
| 60 | # diffusion settings |
| 61 | schedule_type="cosine", |
| 62 | n_steps=3_000, |
| 63 | max_iterations=400_000, |
| 64 | # device setting |
| 65 | device="cuda:0", |
| 66 | # optimizer settings |
| 67 | batch_size=128, |
| 68 | lr_d=1e-4, |
| 69 | show_recon=False, |
| 70 | # pretrain settings |
| 71 | pretrain_weight=None, |
| 72 | pretrain_iterations=None, |
| 73 | *, |
| 74 | # just for debugging |
| 75 | constrain_channel=None, |
| 76 | ): |
| 77 | """train and valid function |
| 78 | |
| 79 | Args: |
| 80 | train_dataset_path (str): _description_ |
| 81 | valid_dataset_path (str): _description_ |
| 82 | batch_size (int, optional): _description_. Defaults to 240. |
| 83 | n_steps (int, optional): _description_. Defaults to 1500. |
| 84 | epochs (int, optional): _description_. Defaults to None. |
| 85 | device (str, optional): _description_. Defaults to 'cuda:0'. |
| 86 | max_iterations (int, optional): _description_. Defaults to 500_000. |
| 87 | lr (float, optional): _description_. Defaults to 1e-4. |
| 88 | pretrain_weight (str, optional): _description_. Defaults to None. |
| 89 | recon_loss (bool, optional): _description_. Defaults to False. |
| 90 | show_recon (bool, optional): _description_. Defaults to False. |
| 91 | constrain_channel (int, optional): _description_. Defaults to None. |
| 92 | """ |
| 93 | from diffusion.diffusion_ddpm_pan import GaussianDiffusion |
| 94 | from models.sr3_dwt import UNetSR3 as Unet |
| 95 | |
| 96 | # init logger |
| 97 | stf_time = time.strftime("%m-%d_%H-%M", time.localtime()) |
| 98 | comment = "pandiff" |
| 99 | logger = TensorboardLogger(file_logger_name="{}-{}".format(stf_time, comment)) |
| 100 | |
| 101 | dataset_name = ( |
| 102 | train_dataset_path.strip(".h5").split("_")[-1] |
| 103 | if not exist(dataset_name) |
| 104 | else dataset_name |
| 105 | ) |
| 106 | logger.print(f"dataset name: {dataset_name}") |
| 107 | division_dict = {"wv3": 2047.0, "gf2": 1023.0, "qb": 2047.0, "cave": 1.0, "harvard": 1.0} |
| 108 | logger.print(f"dataset norm division: {division_dict[dataset_name]}") |
| 109 | rgb_channel = { |
no test coverage detected