else:
optimizer = optim.Adam(model_pos_train.parameters(), lr=lr, amsgrad=True)
lr_decay = args.lr_decay
losses_3d_train = []
losses_3d_train_eval = []
losses_3d_valid = []
epoch = 0
initial_momentum = 0.1
final_momentum = 0.001
train_generator = ChunkedGenerator(args.batch_size//args.stride, cameras_train, poses_train, poses_train_2d, args.stride,
pad=pad, causal_shift=causal_shift, shuffle=True, augment=args.data_augmentation,
kps_left=kps_left, kps_right=kps_right, joints_left=joints_left, joints_right=joints_right)
train_generator_eval = UnchunkedGenerator(cameras_train, poses_train, poses_train_2d,
pad=pad, causal_shift=causal_shift, augment=False)
print('INFO: Training on {} frames'.format(train_generator_eval.num_frames()))
if semi_supervised:
semi_generator = ChunkedGenerator(args.batch_size//args.stride, cameras_semi, None, poses_semi_2d, args.stride,
pad=pad, causal_shift=causal_shift, shuffle=True,
random_seed=4321, augment=args.data_augmentation,
kps_left=kps_left, kps_right=kps_right, joints_left=joints_left, joints_right=joints_right,
endless=True)
semi_generator_eval = UnchunkedGenerator(cameras_semi, None, poses_semi_2d,
pad=pad, causal_shift=causal_shift, augment=False)
print('INFO: Semi-supervision on {} frames'.format(semi_generator_eval.num_frames()))
if args.resume:
optimizer = optim.Adam(model_pos_train.parameters(), lr=lr, amsgrad=True)
lr_decay = args.lr_decay
losses_3d_train = []
losses_3d_train_eval = []
losses_3d_valid = []
epoch = 0
initial_momentum = 0.1
final_momentum = 0.001
train_generator = ChunkedGenerator(args.batch_size//args.stride, cameras_train, poses_train, poses_train_2d, args.randnum, boneindex, args.augdegree, args.stride,
pad=pad, causal_shift=causal_shift, shuffle=True, augment=args.data_augmentation,
kps_left=kps_left, kps_right=kps_right, joints_left=joints_left, joints_right=joints_right)
train_generator_eval = UnchunkedGenerator(cameras_train, poses_train, poses_train_2d,
pad=pad, causal_shift=causal_shift, augment=False)
print('INFO: Training on {} frames'.format(train_generator_eval.num_frames()))
if args.resume:
epoch = checkpoint['epoch']
if 'optimizer' in checkpoint and checkpoint['optimizer'] is not None:
optimizer.load_state_dict(checkpoint['optimizer'])
train_generator.set_random_state(checkpoint['random_state'])
else:
print('WARNING: this checkpoint does not contain an optimizer state. The optimizer will be reinitialized.')
lr = checkpoint['lr']
final_momentum = 0.001
print('3d', poses_train)
print('2d', poses_train_2d)
print('3d shape', np.shape(poses_train))
print('2d shape', np.shape(poses_train_2d))
print(poses_train)
train_generator = ChunkedGenerator(args.batch_size // args.stride,
cameras_train,
poses_train,
poses_train_2d,
args.stride,
pad=pad,
causal_shift=causal_shift,
shuffle=True,
augment=args.data_augmentation,
kps_left=kps_left,
kps_right=kps_right,
joints_left=joints_left,
joints_right=joints_right)
train_generator_eval = UnchunkedGenerator(cameras_train,
poses_train,
poses_train_2d,
pad=pad,
causal_shift=causal_shift,
augment=False)
print('INFO: Training on {} frames'.format(
train_generator_eval.num_frames()))
if args.resume or args.evaluate:
chk_filename = os.path.join(args.checkpoint,
args.resume if args.resume else args.evaluate)
print('Loading checkpoint', chk_filename)
checkpoint = torch.load(chk_filename,
map_location=lambda storage, loc: storage)
model_pos_train.load_state_dict(checkpoint['model_pos'])
model_pos.load_state_dict(checkpoint['model_pos'])
test_generator = ChunkedGenerator(args.batch_size // args.stride,
cameras_valid,
poses_valid,
poses_valid_2d,
args.stride,
pad=pad,
causal_shift=causal_shift,
augment=args.test_time_augmentation,
shuffle=False,
kps_left=kps_left,
kps_right=kps_right,
joints_left=joints_left,
joints_right=joints_right,
noisy=False)
print('INFO: Testing on {} sequences'.format(test_generator.num_frames()))
if not args.evaluate:
cameras_train, poses_train, poses_train_2d = fetch(subjects_train,
action_filter,
subset=args.subset)
lr = args.learning_rate
optimizer = Ranger(model_pos_train.parameters(), lr=lr)
