def main():
FLAGS = parse_args()
cfg = load_config(FLAGS.config)
# cfg['output_eval'] = FLAGS.output_eval
merge_config(FLAGS.opt)
if cfg.use_gpu:
paddle.set_device('gpu')
else:
paddle.set_device('cpu')
if 'slim' in cfg:
cfg = build_slim_model(cfg, mode='test')
check_config(cfg)
check_gpu(cfg.use_gpu)
check_version()
# build trainer
trainer = Trainer(cfg, mode='test')
# load weights
trainer.load_weights(cfg.weights)
# export model
trainer.export(output_dir=FLAGS.save_inference_dir)
def run(FLAGS, cfg):
init_parallel_env()
if FLAGS.enable_ce:
set_random_seed(0)
# build trainer
trainer = Trainer(cfg, mode='train')
# load weights
if 'pretrain_weights' in cfg and cfg.pretrain_weights:
trainer.load_weights(cfg.pretrain_weights)
# init config
if FLAGS.distill_config is not None:
distill_config = load_config(FLAGS.distill_config)
trainer.distill_model = build_teacher_model(distill_config)
trainer.distill_loss = build_distill_loss(distill_config)
trainer.init_optimizer()
# training
trainer.train(FLAGS.eval)
def run(FLAGS, cfg):
# init parallel environment if nranks > 1
init_parallel_env()
# build trainer
trainer = Trainer(cfg, mode='eval')
# load weights
trainer.load_weights(cfg.weights)
# training
trainer.evaluate()
def run(FLAGS, cfg):
# build trainer
trainer = Trainer(cfg, mode='test')
# load weights
trainer.load_weights(cfg.weights)
# get inference images
images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)
# inference
trainer.predict(images,
draw_threshold=FLAGS.draw_threshold,
output_dir=FLAGS.output_dir,
save_txt=FLAGS.save_txt)
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()
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 = TCMRLoss(
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 ========= #
generator = TCMR(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).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 ========= #
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()