Spaces:
Runtime error
Runtime error
| """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) | |
| # dataset parameters | |
| # for train | |
| 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" | |
| ) | |
| # for val | |
| 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) | |
| # visualization parameters | |
| 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" | |
| ) | |
| # network saving and loading parameters | |
| 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") | |
| # training parameters | |
| 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 | |