def test_samples(self):
''' given the same network weights, the random pose generator must produce the same pose for a seed'''
ps = Configer(default_ps_fname='../human_body_prior/train/vposer_smpl_defaults.ini')
vposer = VPoser(num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape = ps.data_shape)
body_pose_rnd = vposer.sample_poses(num_poses=1, seed=100)
body_pose_gt = np.load('samples/body_pose_rnd.npz')['data']
self.assertAlmostEqual(np.square((c2c(body_pose_rnd) - body_pose_gt)).sum(), 0.0)
def __init__(self, work_dir, ps):
from tensorboardX import SummaryWriter
from human_body_prior.data.dataloader import VPoserDS
self.pt_dtype = torch.float64 if ps.fp_precision == '64' else torch.float32
torch.manual_seed(ps.seed)
ps.work_dir = makepath(work_dir, isfile=False)
logger = log2file(os.path.join(work_dir, '%s.log' % ps.expr_code))
summary_logdir = os.path.join(work_dir, 'summaries')
self.swriter = SummaryWriter(log_dir=summary_logdir)
logger('tensorboard --logdir=%s' % summary_logdir)
logger('Torch Version: %s\n' % torch.__version__)
shutil.copy2(os.path.realpath(__file__), work_dir)
use_cuda = torch.cuda.is_available()
if use_cuda: torch.cuda.empty_cache()
self.comp_device = torch.device(
"cuda:%d" % ps.cuda_id if torch.cuda.is_available() else "cpu")
logger('%d CUDAs available!' % torch.cuda.device_count())
gpu_brand = torch.cuda.get_device_name(
ps.cuda_id) if use_cuda else None
logger('Training with %s [%s]' %
(self.comp_device,
gpu_brand) if use_cuda else 'Training on CPU!!!')
logger('Base dataset_dir is %s' % ps.dataset_dir)
kwargs = {'num_workers': ps.n_workers}
ds_train = VPoserDS(dataset_dir=os.path.join(ps.dataset_dir, 'train'))
self.ds_train = DataLoader(ds_train,
batch_size=ps.batch_size,
shuffle=True,
drop_last=True,
**kwargs)
ds_val = VPoserDS(dataset_dir=os.path.join(ps.dataset_dir, 'vald'))
self.ds_val = DataLoader(ds_val,
batch_size=ps.batch_size,
shuffle=True,
drop_last=True,
**kwargs)
ds_test = VPoserDS(dataset_dir=os.path.join(ps.dataset_dir, 'test'))
self.ds_test = DataLoader(ds_test,
batch_size=ps.batch_size,
shuffle=True,
drop_last=True,
**kwargs)
logger('Train dataset size %.2f M' %
(len(self.ds_train.dataset) * 1e-6))
logger('Validation dataset size %d' % len(self.ds_val.dataset))
logger('Test dataset size %d' % len(self.ds_test.dataset))
ps.data_shape = list(ds_val[0]['pose_aa'].shape)
self.vposer_model = VPoser(num_neurons=ps.num_neurons,
latentD=ps.latentD,
data_shape=ps.data_shape,
use_cont_repr=ps.use_cont_repr)
if ps.use_multigpu:
self.vposer_model = nn.DataParallel(self.vposer_model)
self.vposer_model.to(self.comp_device)
varlist = [var[1] for var in self.vposer_model.named_parameters()]
params_count = sum(p.numel() for p in varlist if p.requires_grad)
logger('Total Trainable Parameters Count is %2.2f M.' %
((params_count) * 1e-6))
self.optimizer = optim.Adam(varlist,
lr=ps.base_lr,
weight_decay=ps.reg_coef)
self.logger = logger
self.best_loss_total = np.inf
self.try_num = ps.try_num
self.epochs_completed = 0
self.ps = ps
if ps.best_model_fname is not None:
if isinstance(self.vposer_model, torch.nn.DataParallel):
self.vposer_model.module.load_state_dict(
torch.load(ps.best_model_fname,
map_location=self.comp_device))
else:
self.vposer_model.load_state_dict(
torch.load(ps.best_model_fname,
map_location=self.comp_device))
logger('Restored model from %s' % ps.best_model_fname)
chose_ids = np.random.choice(list(range(len(ds_val))),
size=ps.num_bodies_to_display,
replace=False,
p=None)
data_all = {}
for id in chose_ids:
for k, v in ds_val[id].items():
if k in data_all.keys():
data_all[k] = torch.cat([data_all[k], v[np.newaxis]],
dim=0)
else:
data_all[k] = v[np.newaxis]
self.vis_dorig = {
k: data_all[k].to(self.comp_device)
for k in data_all.keys()
}
self.bm = BodyModel(self.ps.bm_path,
'smplh',
batch_size=self.ps.batch_size,
use_posedirs=True).to(self.comp_device)