def main(cfg):
eval_name = cfg.TRAIN.get('DATASET_EVAL')
print('...Evaluating on {} test set...'.format(eval_name))
vibe = VIBE(
n_layers=cfg.MODEL.TGRU.NUM_LAYERS,
batch_size=cfg.TRAIN.BATCH_SIZE,
seqlen=cfg.DATASET.SEQLEN,
hidden_size=cfg.MODEL.TGRU.HIDDEN_SIZE,
pretrained=cfg.TRAIN.PRETRAINED_REGRESSOR,
add_linear=cfg.MODEL.TGRU.ADD_LINEAR,
bidirectional=cfg.MODEL.TGRU.BIDIRECTIONAL,
use_residual=cfg.MODEL.TGRU.RESIDUAL,
).to(cfg.DEVICE)
if cfg.TRAIN.PRETRAINED_VIBE != '' and os.path.isfile(
cfg.TRAIN.PRETRAINED_VIBE):
checkpoint = torch.load(cfg.TRAIN.PRETRAINED_VIBE)
best_performance = checkpoint['performance']
vibe.load_state_dict(checkpoint['gen_state_dict'])
print(
f'==> Loaded pretrained model from {cfg.TRAIN.PRETRAINED_VIBE}...')
print(f'Performance on 3DPW test set {best_performance}')
else:
print(f'{cfg.TRAIN.PRETRAINED_VIBE} is not a pretrained model!!!!')
model = REFINER(vibe=vibe).to(cfg.DEVICE)
if cfg.TRAIN.PRETRAINED != '' and os.path.isfile(cfg.TRAIN.PRETRAINED):
checkpoint = torch.load(cfg.TRAIN.PRETRAINED)
best_performance = checkpoint['performance']
model.load_state_dict(checkpoint['gen_state_dict'])
print(f'==> Loaded pretrained model from {cfg.TRAIN.PRETRAINED}...')
print(f'Performance on 3DPW test set {best_performance}')
else:
print(f'{cfg.TRAIN.PRETRAINED} is not a pretrained model!!!!')
exit()
test_db = eval(eval_name)(set='test',
seqlen=cfg.DATASET.SEQLEN,
debug=cfg.DEBUG)
test_loader = DataLoader(
dataset=test_db,
batch_size=cfg.TRAIN.BATCH_SIZE,
shuffle=False,
num_workers=cfg.NUM_WORKERS,
)
# reinder_cfg = Config("zen_rf_9",local=True)
# refiner = get_smpl_refiner(reinder_cfg, cfg.DEVICE, num_vae_epoch= 20)
Evaluator(
model=model,
device=cfg.DEVICE,
test_loader=test_loader,
).run()
def main(cfg):
print('...Evaluating on 3DPW test set...')
model = VIBE(
n_layers=cfg.MODEL.TGRU.NUM_LAYERS,
batch_size=cfg.TRAIN.BATCH_SIZE,
seqlen=cfg.DATASET.SEQLEN,
hidden_size=cfg.MODEL.TGRU.HIDDEN_SIZE,
pretrained=cfg.TRAIN.PRETRAINED_REGRESSOR,
add_linear=cfg.MODEL.TGRU.ADD_LINEAR,
bidirectional=cfg.MODEL.TGRU.BIDIRECTIONAL,
use_residual=cfg.MODEL.TGRU.RESIDUAL,
).to(cfg.DEVICE)
if cfg.TRAIN.PRETRAINED != '' and os.path.isfile(cfg.TRAIN.PRETRAINED):
checkpoint = torch.load(cfg.TRAIN.PRETRAINED,
map_location=torch.device('cpu'))
best_performance = checkpoint['performance']
model.load_state_dict(checkpoint['gen_state_dict'])
print(f'==> Loaded pretrained model from {cfg.TRAIN.PRETRAINED}...')
print(f'Performance on 3DPW test set {best_performance}')
else:
print(f'{cfg.TRAIN.PRETRAINED} is not a pretrained model!!!!')
exit()
test_db = ThreeDPW(set='test', seqlen=cfg.DATASET.SEQLEN, debug=cfg.DEBUG)
test_loader = DataLoader(
dataset=test_db,
batch_size=cfg.TRAIN.BATCH_SIZE,
shuffle=False,
num_workers=cfg.NUM_WORKERS,
)
Evaluator(
model=model,
device=cfg.DEVICE,
test_loader=test_loader,
).run()
def main(cfg):
if cfg.SEED_VALUE >= 0:
print(f'Seed value for the experiment {cfg.SEED_VALUE}')
os.environ['PYTHONHASHSEED'] = str(cfg.SEED_VALUE)
random.seed(cfg.SEED_VALUE)
torch.manual_seed(cfg.SEED_VALUE)
np.random.seed(cfg.SEED_VALUE)
logger = create_logger(cfg.LOGDIR, phase='train')
logger.info(f'GPU name -> {torch.cuda.get_device_name()}')
logger.info(f'GPU feat -> {torch.cuda.get_device_properties("cuda")}')
logger.info(pprint.pformat(cfg))
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
writer = SummaryWriter(log_dir=cfg.LOGDIR)
writer.add_text('config', pprint.pformat(cfg), 0)
# ========= Dataloaders ========= #
data_loaders = get_data_loaders(cfg)
# ========= Compile Loss ========= #
loss = VIBELoss(
e_loss_weight=cfg.LOSS.KP_2D_W,
e_3d_loss_weight=cfg.LOSS.KP_3D_W,
e_pose_loss_weight=cfg.LOSS.POSE_W,
e_shape_loss_weight=cfg.LOSS.SHAPE_W,
d_motion_loss_weight=cfg.LOSS.D_MOTION_LOSS_W,
)
# ========= Initialize networks, optimizers and lr_schedulers ========= #
vibe = VIBE(
n_layers=cfg.MODEL.TGRU.NUM_LAYERS,
batch_size=cfg.TRAIN.BATCH_SIZE,
seqlen=cfg.DATASET.SEQLEN,
hidden_size=cfg.MODEL.TGRU.HIDDEN_SIZE,
pretrained=cfg.TRAIN.PRETRAINED_REGRESSOR,
add_linear=cfg.MODEL.TGRU.ADD_LINEAR,
bidirectional=cfg.MODEL.TGRU.BIDIRECTIONAL,
use_residual=cfg.MODEL.TGRU.RESIDUAL,
).to(cfg.DEVICE)
if cfg.TRAIN.PRETRAINED != '' and os.path.isfile(cfg.TRAIN.PRETRAINED):
checkpoint = torch.load(cfg.TRAIN.PRETRAINED)
best_performance = checkpoint['performance']
vibe.load_state_dict(checkpoint['gen_state_dict'])
print(f'==> Loaded pretrained model from {cfg.TRAIN.PRETRAINED}...')
print(f'Performance on 3DPW test set {best_performance}')
else:
print(f'{cfg.TRAIN.PRETRAINED} is not a pretrained model!!!!')
generator = REFINER(vibe=vibe).to(cfg.DEVICE)
gen_optimizer = get_optimizer(
model=generator,
optim_type=cfg.TRAIN.GEN_OPTIM,
lr=cfg.TRAIN.GEN_LR,
weight_decay=cfg.TRAIN.GEN_WD,
momentum=cfg.TRAIN.GEN_MOMENTUM,
)
# motion_discriminator = MotionDiscriminator(
# rnn_size=cfg.TRAIN.MOT_DISCR.HIDDEN_SIZE,
# input_size=69,
# num_layers=cfg.TRAIN.MOT_DISCR.NUM_LAYERS,
# output_size=1,
# feature_pool=cfg.TRAIN.MOT_DISCR.FEATURE_POOL,
# attention_size=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL !='attention' else cfg.TRAIN.MOT_DISCR.ATT.SIZE,
# attention_layers=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL !='attention' else cfg.TRAIN.MOT_DISCR.ATT.LAYERS,
# attention_dropout=None if cfg.TRAIN.MOT_DISCR.FEATURE_POOL !='attention' else cfg.TRAIN.MOT_DISCR.ATT.DROPOUT
# ).to(cfg.DEVICE)
# dis_motion_optimizer = get_optimizer(
# model=motion_discriminator,
# optim_type=cfg.TRAIN.MOT_DISCR.OPTIM,
# lr=cfg.TRAIN.MOT_DISCR.LR,
# weight_decay=cfg.TRAIN.MOT_DISCR.WD,
# momentum=cfg.TRAIN.MOT_DISCR.MOMENTUM
# )
# motion_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
# dis_motion_optimizer,
# mode='min',
# factor=0.1,
# patience=cfg.TRAIN.LR_PATIENCE,
# verbose=True,
# )
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
gen_optimizer,
mode='min',
factor=0.1,
patience=cfg.TRAIN.LR_PATIENCE,
verbose=True,
)
# ========= Start Training ========= #
motion_discriminator = None
dis_motion_optimizer = None
motion_lr_scheduler = None
Trainer(
data_loaders=data_loaders,
generator=generator,
motion_discriminator=motion_discriminator,
criterion=loss,
dis_motion_optimizer=dis_motion_optimizer,
dis_motion_update_steps=cfg.TRAIN.MOT_DISCR.UPDATE_STEPS,
gen_optimizer=gen_optimizer,
start_epoch=cfg.TRAIN.START_EPOCH,
end_epoch=cfg.TRAIN.END_EPOCH,
device=cfg.DEVICE,
writer=writer,
debug=cfg.DEBUG,
logdir=cfg.LOGDIR,
lr_scheduler=lr_scheduler,
motion_lr_scheduler=motion_lr_scheduler,
resume=cfg.TRAIN.RESUME,
num_iters_per_epoch=cfg.TRAIN.NUM_ITERS_PER_EPOCH,
debug_freq=cfg.DEBUG_FREQ,
).fit()
from lib.core.config import VIBE_DATA_DIR, VIBE_DB_DIR
from lib.data_utils.kp_utils import convert_kps
from lib.data_utils.img_utils import normalize_2d_kp, split_into_chunks, transfrom_keypoints
from copycat.smpllib.smpl_mujoco import SMPL_M_Renderer
from zen_renderer.utils.transform_utils import vertizalize_smpl_root
from zen_renderer.utils.image_utils import assemble_videos
if __name__ == "__main__":
device = (torch.device("cuda", index=0)
if torch.cuda.is_available() else torch.device("cpu"))
vibe_model = VIBE(
n_layers=2,
batch_size=32,
seqlen=16,
hidden_size=1024,
add_linear=True,
bidirectional=False,
use_residual=True,
).to(device)
vibe_dir = 'data/vibe_data/vibe_model_wo_3dpw.pth.tar'
checkpoint = torch.load(vibe_dir)
best_performance = checkpoint['performance']
vibe_model.load_state_dict(checkpoint['gen_state_dict'])
vibe_model.eval()
print(f'==> Loaded pretrained model from {vibe_dir}...')
print(f'Performance on 3DPW test set {best_performance}')
meva_model = MEVAV2(90,
hidden_size=1024,
pred_vertices.shape[0], -1, -1).to(pred_vertices.device)
pred_joints = torch.matmul(J_regressor_batch, pred_vertices)
pred_joints = pred_joints[:, H36M_TO_J14, :]
return pred_joints
if __name__ == "__main__":
cfg, cfg_file = parse_args()
eval_name = cfg.TRAIN.get('DATASET_EVAL')
print('...Evaluating on {} test set...'.format(eval_name))
model = VIBE(
n_layers=cfg.MODEL.TGRU.NUM_LAYERS,
batch_size=cfg.TRAIN.BATCH_SIZE,
seqlen=cfg.DATASET.SEQLEN,
hidden_size=cfg.MODEL.TGRU.HIDDEN_SIZE,
pretrained=cfg.TRAIN.PRETRAINED_REGRESSOR,
add_linear=cfg.MODEL.TGRU.ADD_LINEAR,
bidirectional=cfg.MODEL.TGRU.BIDIRECTIONAL,
use_residual=cfg.MODEL.TGRU.RESIDUAL,
).to(cfg.DEVICE)
if cfg.TRAIN.PRETRAINED != '' and os.path.isfile(cfg.TRAIN.PRETRAINED):
checkpoint = torch.load(cfg.TRAIN.PRETRAINED)
best_performance = checkpoint['performance']
model.load_state_dict(checkpoint['gen_state_dict'])
print(f'==> Loaded pretrained model from {cfg.TRAIN.PRETRAINED}...')
print(f'Performance on 3DPW test set {best_performance}')
else:
print(f'{cfg.TRAIN.PRETRAINED} is not a pretrained model!!!!')
exit()