|
import os |
|
import argparse |
|
|
|
def get_parser(): |
|
parser = argparse.ArgumentParser(description="Training config for FontDiffuser.") |
|
|
|
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***.")) |
|
|
|
|
|
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.",) |
|
|
|
|
|
parser.add_argument("--train_batch_size", type=int, default=4, |
|
help="Batch size (per device) for the training dataloader.") |
|
|
|
parser.add_argument("--perceptual_coefficient", type=float, default=0.01) |
|
parser.add_argument("--offset_coefficient", type=float, default=0.5) |
|
|
|
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.") |
|
|
|
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.") |
|
|
|
parser.add_argument("--drop_prob", type=float, default=0.1, help="The uncondition training drop out probability.") |
|
|
|
parser.add_argument("--beta_scheduler", type=str, default="scaled_linear", help="The beta scheduler for DDPM.") |
|
|
|
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.") |
|
|
|
|
|
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 |
|
|