import argparse
from .get_opt import get_opt
from os.path import join as pjoin
import os
class TrainOptions():
def __init__(self):
self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
self.initialized = False
def initialize(self):
# base set
self.parser.add_argument('--name', type=str, default="test", help='Name of this trial')
self.parser.add_argument('--dataset_name', type=str, default='t2m', help='Dataset Name')
self.parser.add_argument('--feat_bias', type=float, default=5, help='Scales for global motion features and foot contact')
self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
self.parser.add_argument('--log_every', type=int, default=5, help='Frequency of printing training progress (by iteration)')
self.parser.add_argument('--save_interval', type=int, default=10_000, help='Frequency of evaluateing and saving models (by iteration)')
# network hyperparams
self.parser.add_argument('--num_layers', type=int, default=8, help='num_layers of transformer')
self.parser.add_argument('--latent_dim', type=int, default=512, help='latent_dim of transformer')
self.parser.add_argument('--text_latent_dim', type=int, default=256, help='latent_dim of text embeding')
self.parser.add_argument('--time_dim', type=int, default=512, help='latent_dim of timesteps')
self.parser.add_argument('--base_dim', type=int, default=512, help='Dimension of Unet base channel')
self.parser.add_argument('--dim_mults', type=int, default=[2,2,2,2], nargs='+', help='Unet channel multipliers.')
self.parser.add_argument('--no_eff', action='store_true', help='whether use efficient linear attention')
self.parser.add_argument('--no_adagn', action='store_true', help='whether use adagn block')
self.parser.add_argument('--diffusion_steps', type=int, default=1000, help='diffusion_steps of transformer')
self.parser.add_argument('--prediction_type', type=str, default='sample', help='diffusion_steps of transformer')
# train hyperparams
self.parser.add_argument('--seed', type=int, default=0, help='seed for train')
self.parser.add_argument('--num_train_steps', type=int, default=50_000, help='Number of training iterations')
self.parser.add_argument('--lr', type=float, default=2e-4, help='Learning rate')
self.parser.add_argument("--decay_rate", default=0.9, type=float, help="the decay rate of lr (0-1 default 0.9)")
self.parser.add_argument("--update_lr_steps", default=5_000, type=int, help="")
self.parser.add_argument("--cond_mask_prob", default=0.1, type=float,
help="The probability of masking the condition during training."
" For classifier-free guidance learning.")
self.parser.add_argument('--clip_grad_norm', type=float, default=1, help='Gradient clip')
self.parser.add_argument('--weight_decay', type=float, default=1e-2, help='Learning rate weight_decay')
self.parser.add_argument('--batch_size', type=int, default=64, help='Batch size per GPU')
self.parser.add_argument("--beta_schedule", default='linear', type=str, help="Types of beta in diffusion (e.g. linear, cosine)")
self.parser.add_argument('--dropout', type=float, default=0.1, help='dropout')
# continue training
self.parser.add_argument('--is_continue', action="store_true", help='Is this trail continued from previous trail?')
self.parser.add_argument('--continue_ckpt', type=str, default="latest.tar", help='previous trail to continue')
self.parser.add_argument("--opt_path", type=str, default='',help='option file path for loading model')
self.parser.add_argument('--debug', action="store_true", help='debug mode')
self.parser.add_argument('--self_attention', action="store_true", help='self_attention use or not')
self.parser.add_argument('--vis_attn', action='store_true', help='vis attention value or not')
self.parser.add_argument('--edit_mode', action='store_true', help='editing mode')
# EMA params
self.parser.add_argument(
"--model-ema", action="store_true", help="enable tracking Exponential Moving Average of model parameters"
)
self.parser.add_argument(
"--model-ema-steps",
type=int,
default=32,
help="the number of iterations that controls how often to update the EMA model (default: 32)",
)
self.parser.add_argument(
"--model-ema-decay",
type=float,
default=0.9999,
help="decay factor for Exponential Moving Average of model parameters (default: 0.99988)",
)
self.initialized = True
def parse(self,accelerator):
if not self.initialized:
self.initialize()
self.opt = self.parser.parse_args()
if self.opt.is_continue:
assert self.opt.opt_path.endswith('.txt')
get_opt(self.opt, self.opt.opt_path)
self.opt.is_train = True
self.opt.is_continue=True
elif accelerator.is_main_process:
args = vars(self.opt)
accelerator.print('------------ Options -------------')
for k, v in sorted(args.items()):
accelerator.print('%s: %s' % (str(k), str(v)))
accelerator.print('-------------- End ----------------')
# save to the disk
expr_dir = pjoin(self.opt.checkpoints_dir, self.opt.dataset_name, self.opt.name)
os.makedirs(expr_dir,exist_ok=True)
file_name = pjoin(expr_dir, 'opt.txt')
with open(file_name, 'wt') as opt_file:
opt_file.write('------------ Options -------------\n')
for k, v in sorted(args.items()):
if k =='opt_path':
continue
opt_file.write('%s: %s\n' % (str(k), str(v)))
opt_file.write('-------------- End ----------------\n')
if self.opt.dataset_name == 't2m' or self.opt.dataset_name == 'humanml':
self.opt.joints_num = 22
self.opt.dim_pose = 263
self.opt.max_motion_length = 196
self.opt.radius = 4
self.opt.fps = 20
elif self.opt.dataset_name == 'kit':
self.opt.joints_num = 21
self.opt.dim_pose = 251
self.opt.max_motion_length = 196
self.opt.radius = 240 * 8
self.opt.fps = 12.5
else:
raise KeyError('Dataset not recognized')
self.opt.device = accelerator.device
self.opt.is_train = True
return self.opt