support CMU Mocap
This commit is contained in:
@ -42,10 +42,10 @@ def init(smpl_layer, target, device, cfg):
|
||||
params["scale"] = params["scale"].to(device)
|
||||
|
||||
params["pose_params"].requires_grad = True
|
||||
params["shape_params"].requires_grad = True
|
||||
params["scale"].requires_grad = False
|
||||
params["shape_params"].requires_grad = bool(cfg.TRAIN.OPTIMIZE_SHAPE)
|
||||
params["scale"].requires_grad = bool(cfg.TRAIN.OPTIMIZE_SCALE)
|
||||
|
||||
optimizer = optim.Adam([params["pose_params"], params["shape_params"]],
|
||||
optimizer = optim.Adam([params["pose_params"], params["shape_params"], params["scale"]],
|
||||
lr=cfg.TRAIN.LEARNING_RATE)
|
||||
|
||||
index={}
|
||||
@ -73,8 +73,9 @@ def train(smpl_layer, target,
|
||||
|
||||
early_stop = Early_Stop()
|
||||
for epoch in tqdm(range(cfg.TRAIN.MAX_EPOCH)):
|
||||
# for epoch in range(cfg.TRAIN.MAX_EPOCH):
|
||||
verts, Jtr = smpl_layer(pose_params, th_betas=shape_params)
|
||||
loss = F.smooth_l1_loss(Jtr.index_select(1, index["smpl_index"]) * 100,
|
||||
loss = F.smooth_l1_loss(Jtr.index_select(1, index["smpl_index"]) * 100 * scale,
|
||||
target.index_select(1, index["dataset_index"]) * 100)
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
@ -88,8 +89,8 @@ def train(smpl_layer, target,
|
||||
break
|
||||
|
||||
if epoch % cfg.TRAIN.WRITE == 0:
|
||||
# logger.info("Epoch {}, lossPerBatch={:.6f}, EarlyStopSatis: {}".format(
|
||||
# epoch, float(loss), early_stop.satis_num))
|
||||
# logger.info("Epoch {}, lossPerBatch={:.6f}, scale={:.4f} EarlyStopSatis: {}".format(
|
||||
# epoch, float(loss),float(scale), early_stop.satis_num))
|
||||
writer.add_scalar('loss', float(loss), epoch)
|
||||
writer.add_scalar('learning_rate', float(
|
||||
optimizer.state_dict()['param_groups'][0]['lr']), epoch)
|
||||
|
||||
Reference in New Issue
Block a user