Spaces:
Running
Running
File size: 4,341 Bytes
12deb01 |
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 |
import argparse
import os
import torch
from mmcv.runner import get_dist_info
import torch.distributed as dist
class BaseOptions():
def __init__(self):
self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
self.initialized = False
def initialize(self):
self.parser.add_argument('--name', type=str, default="test", help='Name of this trial')
self.parser.add_argument('--decomp_name', type=str, default="Decomp_SP001_SM001_H512", help='Name of autoencoder model')
self.parser.add_argument("--gpu_id", type=int, default=-1, help='GPU id')
self.parser.add_argument("--distributed", action="store_true", help='Weather to use DDP training')
self.parser.add_argument('--dataset_name', type=str, default='t2m', help='Dataset Name')
self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
self.parser.add_argument("--unit_length", type=int, default=4, help="Motions are cropped to the maximum times of unit_length")
self.parser.add_argument("--max_text_len", type=int, default=20, help="Maximum length of text description")
self.parser.add_argument('--text_enc_mod', type=str, default='bigru')
self.parser.add_argument('--estimator_mod', type=str, default='bigru')
self.parser.add_argument('--dim_text_hidden', type=int, default=512, help='Dimension of hidden unit in text encoder')
self.parser.add_argument('--dim_att_vec', type=int, default=512, help='Dimension of attention vector')
self.parser.add_argument('--dim_z', type=int, default=128, help='Dimension of latent Gaussian vector')
self.parser.add_argument('--n_layers_pri', type=int, default=1, help='Number of layers in prior network')
self.parser.add_argument('--n_layers_pos', type=int, default=1, help='Number of layers in posterior network')
self.parser.add_argument('--n_layers_dec', type=int, default=1, help='Number of layers in generator')
self.parser.add_argument('--dim_pri_hidden', type=int, default=1024, help='Dimension of hidden unit in prior network')
self.parser.add_argument('--dim_pos_hidden', type=int, default=1024, help='Dimension of hidden unit in posterior network')
self.parser.add_argument('--dim_dec_hidden', type=int, default=1024, help='Dimension of hidden unit in generator')
self.parser.add_argument('--dim_movement_enc_hidden', type=int, default=512,
help='Dimension of hidden in AutoEncoder(encoder)')
self.parser.add_argument('--dim_movement_dec_hidden', type=int, default=512,
help='Dimension of hidden in AutoEncoder(decoder)')
self.parser.add_argument('--dim_movement_latent', type=int, default=512, help='Dimension of motion snippet')
self.initialized = True
def parse(self):
if not self.initialized:
self.initialize()
self.opt = self.parser.parse_args()
self.opt.is_train = self.is_train
if self.opt.gpu_id != -1:
# self.opt.gpu_id = int(self.opt.gpu_id)
torch.cuda.set_device(self.opt.gpu_id)
args = vars(self.opt)
if args["distributed"]:
init_dist('slurm')
rank, world_size = get_dist_info()
if rank == 0:
print('------------ Options -------------')
for k, v in sorted(args.items()):
print('%s: %s' % (str(k), str(v)))
print('-------------- End ----------------')
if self.is_train:
# save to the disk
expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.dataset_name, self.opt.name)
if not os.path.exists(expr_dir):
os.makedirs(expr_dir)
file_name = os.path.join(expr_dir, 'opt.txt')
with open(file_name, 'wt') as opt_file:
opt_file.write('------------ Options -------------\n')
for k, v in sorted(args.items()):
opt_file.write('%s: %s\n' % (str(k), str(v)))
opt_file.write('-------------- End ----------------\n')
if world_size > 1:
dist.barrier()
return self.opt
|