File size: 6,379 Bytes
508b842 |
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 |
import os
import argparse
def get_parser():
parser = argparse.ArgumentParser(description="Training config for FontDiffuser.")
################# Experience #################
parser.add_argument("--seed", type=int, default=123, help="A seed for reproducible training.")
parser.add_argument("--experience_name", type=str, default="fontdiffuer_training")
parser.add_argument("--data_root", type=str, default=None,
help="The font dataset root path.",)
parser.add_argument("--output_dir", type=str, default=None,
help="The output directory where the model predictions and checkpoints will be written.")
parser.add_argument("--report_to", type=str, default="tensorboard")
parser.add_argument("--logging_dir", type=str, default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."))
# Model
parser.add_argument("--resolution", type=int, default=96,
help="The resolution for input images, all the images in the train/validation \
dataset will be resized to this.")
parser.add_argument("--unet_channels", type=tuple, default=(64, 128, 256, 512),
help="The channels of the UNet.")
parser.add_argument("--style_image_size", type=int, default=96, help="The size of style images.")
parser.add_argument("--content_image_size", type=int, default=96, help="The size of content images.")
parser.add_argument("--content_encoder_downsample_size", type=int, default=3,
help="The downsample size of the content encoder.")
parser.add_argument("--channel_attn", type=bool, default=True, help="Whether to use the se attention.",)
parser.add_argument("--content_start_channel", type=int, default=64,
help="The channels of the fisrt layer output of content encoder.",)
parser.add_argument("--style_start_channel", type=int, default=64,
help="The channels of the fisrt layer output of content encoder.",)
# Training
parser.add_argument("--train_batch_size", type=int, default=4,
help="Batch size (per device) for the training dataloader.")
## loss coefficient
parser.add_argument("--perceptual_coefficient", type=float, default=0.01)
parser.add_argument("--offset_coefficient", type=float, default=0.5)
## step
parser.add_argument("--max_train_steps", type=int, default=440000,
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",)
parser.add_argument("--ckpt_interval", type=int,default=40000, help="The step begin to validate.")
parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.",)
parser.add_argument("--log_interval", type=int, default=100, help="The log interval of training.")
## learning rate
parser.add_argument("--learning_rate", type=float, default=1e-4,
help="Initial learning rate (after the potential warmup period) to use.")
parser.add_argument("--scale_lr", action="store_true", default=False,
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
parser.add_argument("--lr_scheduler", type=str, default="linear",
help="The scheduler type to use. Choose between 'linear', 'cosine', \
'cosine_with_restarts', 'polynomial', 'constant', 'constant_with_warmup'")
parser.add_argument("--lr_warmup_steps", type=int, default=10000,
help="Number of steps for the warmup in the lr scheduler.")
## classifier-free
parser.add_argument("--drop_prob", type=float, default=0.1, help="The uncondition training drop out probability.")
## scheduler
parser.add_argument("--beta_scheduler", type=str, default="scaled_linear", help="The beta scheduler for DDPM.")
## optimizer
parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument("--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16"],
help="Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires \
PyTorch >= 1.10. and an Nvidia Ampere GPU.")
# Sampling
parser.add_argument("--algorithm_type", type=str, default="dpmsolver++", help="Algorithm for sampleing.")
parser.add_argument("--guidance_type", type=str, default="classifier-free", help="Guidance type of sampling.")
parser.add_argument("--guidance_scale", type=float, default=7.5, help="Guidance scale of the classifier-free mode.")
parser.add_argument("--num_inference_steps", type=int, default=20, help="Sampling step.")
parser.add_argument("--model_type", type=str, default="noise", help="model_type for sampling.")
parser.add_argument("--order", type=int, default=2, help="The order of the dpmsolver.")
parser.add_argument("--skip_type", type=str, default="time_uniform", help="Skip type of dpmsolver.")
parser.add_argument("--method", type=str, default="multistep", help="Multistep of dpmsolver.")
parser.add_argument("--correcting_x0_fn", type=str, default=None, help="correcting_x0_fn of dpmsolver.")
parser.add_argument("--t_start", type=str, default=None, help="t_start of dpmsolver.")
parser.add_argument("--t_end", type=str, default=None, help="t_end of dpmsolver.")
parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
return parser
|