|
|
"""This script contains the training options for Deep3DFaceRecon_pytorch |
|
|
""" |
|
|
from util import util |
|
|
|
|
|
from .base_options import BaseOptions |
|
|
|
|
|
|
|
|
class TrainOptions(BaseOptions): |
|
|
"""This class includes training options. |
|
|
|
|
|
It also includes shared options defined in BaseOptions. |
|
|
""" |
|
|
|
|
|
def initialize(self, parser): |
|
|
parser = BaseOptions.initialize(self, parser) |
|
|
|
|
|
|
|
|
parser.add_argument("--data_root", type=str, default="./", help="dataset root") |
|
|
parser.add_argument( |
|
|
"--flist", type=str, default="datalist/train/masks.txt", help="list of mask names of training set" |
|
|
) |
|
|
parser.add_argument("--batch_size", type=int, default=32) |
|
|
parser.add_argument( |
|
|
"--dataset_mode", type=str, default="flist", help="chooses how datasets are loaded. [None | flist]" |
|
|
) |
|
|
parser.add_argument( |
|
|
"--serial_batches", |
|
|
action="store_true", |
|
|
help="if true, takes images in order to make batches, otherwise takes them randomly", |
|
|
) |
|
|
parser.add_argument("--num_threads", default=4, type=int, help="# threads for loading data") |
|
|
parser.add_argument( |
|
|
"--max_dataset_size", |
|
|
type=int, |
|
|
default=float("inf"), |
|
|
help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--preprocess", |
|
|
type=str, |
|
|
default="shift_scale_rot_flip", |
|
|
help="scaling and cropping of images at load time [shift_scale_rot_flip | shift_scale | shift | shift_rot_flip ]", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--use_aug", type=util.str2bool, nargs="?", const=True, default=True, help="whether use data augmentation" |
|
|
) |
|
|
|
|
|
|
|
|
parser.add_argument( |
|
|
"--flist_val", type=str, default="datalist/val/masks.txt", help="list of mask names of val set" |
|
|
) |
|
|
parser.add_argument("--batch_size_val", type=int, default=32) |
|
|
|
|
|
|
|
|
parser.add_argument( |
|
|
"--display_freq", type=int, default=1000, help="frequency of showing training results on screen" |
|
|
) |
|
|
parser.add_argument( |
|
|
"--print_freq", type=int, default=100, help="frequency of showing training results on console" |
|
|
) |
|
|
|
|
|
|
|
|
parser.add_argument("--save_latest_freq", type=int, default=5000, help="frequency of saving the latest results") |
|
|
parser.add_argument( |
|
|
"--save_epoch_freq", type=int, default=1, help="frequency of saving checkpoints at the end of epochs" |
|
|
) |
|
|
parser.add_argument("--evaluation_freq", type=int, default=5000, help="evaluation freq") |
|
|
parser.add_argument("--save_by_iter", action="store_true", help="whether saves model by iteration") |
|
|
parser.add_argument("--continue_train", action="store_true", help="continue training: load the latest model") |
|
|
parser.add_argument( |
|
|
"--epoch_count", |
|
|
type=int, |
|
|
default=1, |
|
|
help="the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...", |
|
|
) |
|
|
parser.add_argument("--phase", type=str, default="train", help="train, val, test, etc") |
|
|
parser.add_argument("--pretrained_name", type=str, default=None, help="resume training from another checkpoint") |
|
|
|
|
|
|
|
|
parser.add_argument("--n_epochs", type=int, default=20, help="number of epochs with the initial learning rate") |
|
|
parser.add_argument("--lr", type=float, default=0.0001, help="initial learning rate for adam") |
|
|
parser.add_argument( |
|
|
"--lr_policy", type=str, default="step", help="learning rate policy. [linear | step | plateau | cosine]" |
|
|
) |
|
|
parser.add_argument( |
|
|
"--lr_decay_epochs", type=int, default=10, help="multiply by a gamma every lr_decay_epochs epoches" |
|
|
) |
|
|
|
|
|
self.isTrain = True |
|
|
return parser |
|
|
|
