def run(args): import hyperparameters cfg = hyperparameters.get_config(args) # cfg.layers_per_scale=1 # cfg.num_keypoints=32 # cfg.batch_size = 25 imgs_to_keyp_model = ImagesToKeypEncoder(cfg, (8, 3, 64, 64), debug=True) keyp_to_imgs_model = KeypToImagesDecoder(cfg, (8, 3, 64, 64), debug=True) print(imgs_to_keyp_model) print(keyp_to_imgs_model) #summary(model, input_size=(2, 3, 64, 64)) k, h = imgs_to_keyp_model(0.5 * torch.ones((1, 8, 3, 64, 64))) r = keyp_to_imgs_model(k) print(k.shape, h.shape, r.shape) b = sum([ np.prod(list(params.size())) for params in imgs_to_keyp_model.parameters() ]) print("Encodeer params: ", b) c = sum([ np.prod(list(params.size())) for params in keyp_to_imgs_model.parameters() ]) print("Decoder params: ", c) print("Model parameters: ", b + c)
def run(args): import hyperparameters cfg = hyperparameters.get_config(args) cfg.layers_per_scale = 1 # cfg.layers_per_scale=1 # cfg.num_keypoints=32 # cfg.batch_size = 25 utils.set_seed_everywhere(args.seed) imgs_to_keyp_model = ImagesToKeypEncoder(cfg, (8, 3, 64, 64), debug=True) keyp_to_imgs_model = KeypToImagesDecoder(cfg, (8, 3, 64, 64), debug=True) keyp_pred_net = KeypPredictor(cfg) keyp_inverse_net = KeypInverseModel(cfg) print(imgs_to_keyp_model) print(keyp_to_imgs_model) print(keyp_pred_net) # summary(model, input_size=(2, 3, 64, 64)) img = 0.5 * torch.ones((1, 4, 3, 64, 64)) action = 0.4 * torch.ones((1, 4, 4)) k, h = imgs_to_keyp_model(img) r = keyp_to_imgs_model(k, img[:, 0], k[:, 0]) print(k.shape, h.shape, r.shape) pred_k = keyp_pred_net(k[Ellipsis, :2], action) pred_action = keyp_inverse_net(k[Ellipsis, :2]) print("Pred_k: ", pred_k.shape, "Pred_action:", pred_action.shape) b = sum([ np.prod(list(params.size())) for params in imgs_to_keyp_model.parameters() ]) print("Encodeer params: ", b) c = sum([ np.prod(list(params.size())) for params in keyp_to_imgs_model.parameters() ]) print("Decoder params: ", c) d = sum([ np.prod(list(params.size())) for params in keyp_pred_net.parameters() ]) print("Keyp Predictor params: ", d) print("Model parameters: ", b + c + d) for n in range(k.shape[1]): print(pred_k[0, 2, n, :2], k[0, 2, n, :2]) print( F.mse_loss(pred_k, k[:, 1:, :, :2], reduction='sum') / (pred_k.shape[0] * pred_k.shape[1]))
def main(args): utils.set_seed_everywhere(args.seed) cfg = hyperparameters.get_config(args) cfg.seed = args.seed args.cuda = not args.no_cuda and torch.cuda.is_available() time_str = datetime.now( timezone('US/Eastern')).strftime("%Y-%m-%d-%H-%M-%S") exp_dir = os.path.join(cfg.base_dir, time_str) checkpoint_dir = os.path.join(exp_dir, cfg.checkpoint_dir) log_dir = os.path.join(exp_dir, cfg.log_dir) save_config(cfg, exp_dir, "config.json") print("Log path: ", log_dir, "Checkpoint Dir: ", checkpoint_dir) num_timsteps = cfg.observed_steps + cfg.predicted_steps data_shape = {'image': (None, num_timsteps, 3, 64, 64)} cfg.data_shapes = data_shape model = KeypointModel(cfg) cp_callback = ModelCheckpoint(filepath=os.path.join( checkpoint_dir, "model_"), period=25, save_top_k=-1) logger = TensorBoardLogger(log_dir, name="", version=None) gpus = 1 if args.cuda else None if args.pretrained_path: checkpoint_path = get_latest_checkpoint(args.pretrained_path) import json model = KeypointModel.load_from_checkpoint(checkpoint_path) print(json.dumps(model.cfg, indent=4)) print("On GPU Device: ", gpus) trainer = Trainer( max_epochs=args.num_epochs, logger=logger, checkpoint_callback=cp_callback, gpus=gpus, #distributed_backend='dp', progress_bar_refresh_rate=1, #gradient_clip_val=cfg.clipnorm, fast_dev_run=False, #train_percent_check=0.1,val_percent_check=0.0, #val_percent_check=0.3, track_grad_norm=2, show_progress_bar=True) trainer.fit(model) save_path = os.path.join(checkpoint_dir, "model_final_" + str(args.num_epochs) + ".ckpt") print("Saving model finally:") trainer.save_checkpoint(save_path)
def viz_seq_unroll(args): torch.random.manual_seed(0) np.random.seed(0) cfg = hyperparameters.get_config(args) unroll_T = 16 args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if args.cuda else "cpu") l_dir = cfg.train_dir if args.is_train else args.test_dir print("Data loader: ", l_dir) loader, data_shapes = datasets.get_sequence_dataset( data_dir=os.path.join(cfg.data_dir, l_dir), batch_size=cfg.batch_size, num_timesteps=cfg.observed_steps + cfg.predicted_steps, shuffle=False) cfg.data_shapes = data_shapes model = train_dynamics.KeypointModel(cfg).to(device) if args.pretrained_path: checkpoint_path = get_latest_checkpoint(args.pretrained_path) checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) print("Loading model from: ", checkpoint_path) model.load_state_dict(checkpoint['state_dict']) model.eval() with torch.no_grad(): for data in islice(loader, 1): img_seq = data['image'].to(device) pred_img_seq, pred_keyp_seq = model.unroll(img_seq, unroll_T) bs, T = img_seq.shape[0], img_seq.shape[1] print( "LOSS:", F.mse_loss(img_seq, pred_img_seq[:, :T], reduction='sum') / (bs * T)) print(img_seq.shape, pred_keyp_seq.shape, pred_img_seq.shape) imgs_seq_np, pred_img_seq_np = img_torch_to_numpy( img_seq), img_torch_to_numpy(pred_img_seq) keypoints_seq_np = pred_keyp_seq.cpu().numpy() num_seq = imgs_seq_np.shape[0] for i in islice(range(num_seq), 3): save_path = os.path.join( args.vids_dir, args.vids_path + "_" + l_dir + "_{}.mp4".format(i)) print(i, "Video PRED Save Path", save_path) viz_all_unroll(imgs_seq_np[i], pred_img_seq_np[i], keypoints_seq_np[i], True, 100, save_path)
def main(args): global LOG_INTERVAL, SAVE_PATH torch.manual_seed(args.seed) np.random.seed(args.seed) cfg = hyperparameters.get_config(args) args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if args.cuda else "cpu") print("Using device: ", device) time_str = datetime.now( timezone('US/Eastern')).strftime("%Y-%m-%d-%H-%M-%S") exp_dir = os.path.join(cfg.base_dir, time_str) checkpoint_dir = os.path.join(exp_dir, cfg.checkpoint_dir) log_dir = os.path.join(exp_dir, cfg.log_dir) print("Log path: ", log_dir, "Checkpoint Dir: ", checkpoint_dir) if not os.path.isdir(log_dir): os.makedirs(log_dir) if not os.path.isdir(checkpoint_dir): os.makedirs(checkpoint_dir) LOG_INTERVAL = args.log_interval SAVE_PATH = os.path.join(checkpoint_dir, "model-") save_config(cfg, exp_dir, "config.json") train_loader, data_shapes = datasets.get_dataset(data_dir=os.path.join( cfg.data_dir, cfg.train_dir), batch_size=cfg.batch_size, shuffle=False) test_loader, _ = datasets.get_dataset(data_dir=os.path.join( cfg.data_dir, cfg.test_dir), batch_size=cfg.batch_size, shuffle=False) models = build_model_noseq(cfg, data_shapes).to(device) optimizer = optim.Adam(models.parameters(), lr=cfg.learning_rate, weight_decay=1e-4) model_dict = { 'models': models, 'optimizer': optimizer, 'train_loader': train_loader, 'test_loader': test_loader, 'writer': SummaryWriter(log_dir), 'device': device } for i in tqdm(range(args.num_epochs)): train_epoch(i, model_dict, cfg)
def run_final_test(args): utils.set_seed_everywhere(args.seed) cfg = hyperparameters.get_config(args) args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if args.cuda else "cpu") l_dir = cfg.train_dir if args.is_train else args.test_dir print("Data loader: ", l_dir) loader, data_shapes = datasets.get_sequence_dataset( data_dir=os.path.join(cfg.data_dir, l_dir), batch_size=cfg.batch_size, num_timesteps=2 * args.timesteps, shuffle=True) cfg.log_training = args.log_training cfg.log_training_path = os.path.join(args.log_training_path) cfg.data_shapes = data_shapes if args.no_first: if args.keyp_pred: print("Loding keyp pred") model = train_keyp_pred.KeypointModel(cfg).to(device) elif args.keyp_inverse: print("Loding Inverse Model") model = train_keyp_inverse.KeypointModel(cfg).to(device) else: pass else: model = train.KeypointModel(cfg).to(device) if args.pretrained_path: if args.ckpt: checkpoint_path = os.path.join( args.pretrained_path, "_ckpt_epoch_" + args.ckpt + ".ckpt") else: print("Loading latest") checkpoint_path = get_latest_checkpoint(args.pretrained_path) checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) print("Loading model from: ", checkpoint_path) model.load_state_dict(checkpoint['state_dict']) model.eval() print("Load complete") trainer = Trainer(gpus=1, progress_bar_refresh_rate=1, show_progress_bar=True) trainer.test(model)
def run(args): import hyperparameters torch.set_printoptions(precision=10) cfg = hyperparameters.get_config(args) # cfg.layers_per_scale=1 # cfg.num_keypoints=32 # cfg.batch_size = 25 imgs_to_keyp_model = ImagesToKeypEncoder(cfg, (3, 64, 64), debug=True) imgs_to_keyp_model.apply(weights_init) keyp_to_imgs_model = KeypToImagesDecoder(cfg, (3, 64, 64), debug=True) keyp_to_imgs_model.apply(weights_init) print(imgs_to_keyp_model) print(keyp_to_imgs_model) #summary(model, input_size=(2, 3, 64, 64)) k, h = imgs_to_keyp_model(0.5 * torch.ones((cfg.batch_size, 3, 64, 64))) print(k[0]) k_tf = torch.from_numpy(np.load("../unsup_keyp/tf_test_dec_0.npy")) r, xs = keyp_to_imgs_model(k_tf) for i, y in enumerate(xs): if len(y.shape) == 4: np.save('torch_test_dec_{}.npy'.format(i), y.detach().permute(0, 2, 3, 1).numpy()) else: np.save('torch_test_dec_{}.npy'.format(i), y.detach().numpy()) print(k.shape, h.shape, r.shape) b = sum([ np.prod(list(params.size())) for params in imgs_to_keyp_model.parameters() ]) print("Encodeer params: ", b) c = sum([ np.prod(list(params.size())) for params in keyp_to_imgs_model.parameters() ]) print("Decoder params: ", c) print("Model parameters: ", b + c)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cfg = hyperparameters.get_config(args) args.cuda = not args.no_cuda and torch.cuda.is_available() time_str = datetime.now( timezone('US/Eastern')).strftime("%Y-%m-%d-%H-%M-%S") exp_dir = os.path.join(cfg.base_dir, time_str) checkpoint_dir = os.path.join(exp_dir, cfg.checkpoint_dir) log_dir = os.path.join(exp_dir, cfg.log_dir) save_config(cfg, exp_dir, "config.json") print("Log path: ", log_dir, "Checkpoint Dir: ", checkpoint_dir) data_shape = {'image': (None, 3, 64, 64)} cfg.data_shapes = data_shape model = KeypointModel(cfg) cp_callback = ModelCheckpoint(filepath=os.path.join( checkpoint_dir, "model_"), period=2, save_top_k=-1) logger = TensorBoardLogger(log_dir, name="", version=None) gpus = 1 if args.cuda else None if args.pretrained_path: checkpoint_path = get_latest_checkpoint(args.pretrained_path) import json model = KeypointModel.load_from_checkpoint(checkpoint_path) print(json.dumps(model.cfg, indent=4)) print("On GPU Device: ", gpus) trainer = Trainer(max_epochs=args.num_epochs, logger=logger, checkpoint_callback=cp_callback, gpus=gpus, progress_bar_refresh_rate=1, gradient_clip_val=cfg.clipnorm, fast_dev_run=False) trainer.fit(model)
def main(args): cfg = hyperparameters.get_config(args) args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if args.cuda else "cpu") l_dir = cfg.train_dir if args.is_train else args.test_dir print("Data loader: ", l_dir) loader, data_shapes = datasets.get_dataset(data_dir=os.path.join( cfg.data_dir, l_dir), batch_size=cfg.batch_size) models = build_model_noseq(cfg, data_shapes).to(device) if args.pretrained_path: model_path = get_latest_checkpoint(args.pretrained_path, "*.pth") print("Loading model from: ", model_path) models.load_state_dict(torch.load(model_path)) models.eval() images_to_keypoints_net, keypoints_to_images_net = models with torch.no_grad(): for i, data in islice(enumerate(loader), 5): img = data['image'].to(device) keypoints, _ = images_to_keypoints_net(img) pred_img = keypoints_to_images_net(keypoints) print(img.shape, keypoints.shape, pred_img.shape) save_path = os.path.join( args.vids_dir, args.vids_path + "_" + l_dir + "_{}.mp4".format(i)) print(i, "Video Save Path", save_path) imgs_np, pred_img_np = img_torch_to_numpy(img), img_torch_to_numpy( pred_img) keypoints_np = keypoints.cpu().numpy() viz_all(imgs_np, pred_img_np, keypoints_np, True, 300, save_path)
def load_model(args): utils.set_seed_everywhere(args.seed) cfg = hyperparameters.get_config(args) cfg.data_shapes = {'image': (None, 16, 3, 64, 64)} args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if args.cuda else "cpu") if not args.inv_fwd: model = train_keyp_pred.KeypointModel(cfg).to(device) else: model = train_keyp_inverse_forward.KeypointModel(cfg).to(device) checkpoint_path = os.path.join(args.pretrained_path, "_ckpt_epoch_" + args.ckpt + ".ckpt") checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) print("Loading model from: ", checkpoint_path) model.load_state_dict(checkpoint['state_dict']) model.eval() print("Load complete") return model
import hyperparameters, hyperparameters_tf tf.enable_eager_execution() import numpy as np w_init = 0.00001 def weights_init(m): if isinstance(m, nn.Conv2d): nn.init.constant(m.weight.data, w_init) nn.init.constant(m.bias.data, 0.0) from register_args import get_argparse args = get_argparse(False).parse_args() cfg_tor = hyperparameters.get_config(args) cfg_tf = hyperparameters_tf.get_config(args) x = np.random.randn((2, 8, 64, 64, 3)).astype(np.float32) xtor = torch.from_numpy(x).permute(0, 3, 1, 2) xtf = tf.convert_to_tensor(x) model_tor = KeypointModel(cfg_tor) model_tor.apply(weights_init) model_tf = build_model_keyp(cfg_tf, {'image': (None, 8, 64, 64, 3)})
def viz_seq(args): utils.set_seed_everywhere(args.seed) cfg = hyperparameters.get_config(args) args.cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if args.cuda else "cpu") l_dir = cfg.train_dir if args.is_train else args.test_dir print("Data loader: ", l_dir) loader, data_shapes = datasets.get_sequence_dataset( data_dir=os.path.join(cfg.data_dir, l_dir), batch_size=10, num_timesteps=args.timesteps, shuffle=True) cfg.data_shapes = data_shapes model = train_keyp_inverse_forward.KeypointModel(cfg).to(device) if args.pretrained_path: if args.ckpt: checkpoint_path = os.path.join( args.pretrained_path, "_ckpt_epoch_" + args.ckpt + ".ckpt") else: print("Loading latest") checkpoint_path = get_latest_checkpoint(args.pretrained_path) checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) print("Loading model from: ", checkpoint_path) model.load_state_dict(checkpoint['state_dict']) model.eval() print("Load complete") with torch.no_grad(): for data in islice(loader, 1): img_seq = data['image'].to(device) action_seq = data['action'].to(device) keypoints_seq, heatmaps_seq, pred_keyp_seq, pred_action_seq = model( img_seq, action_seq) print( "Keypoint Pred LOSS:", F.mse_loss(pred_keyp_seq[Ellipsis, :2], keypoints_seq[:, 1:, :, :2], reduction='sum') / ((pred_keyp_seq.shape[0]) * pred_keyp_seq.shape[1])) if args.unroll: pred_keyp_seq = model.unroll(img_seq, action_seq) pred_keyp_seq_np = pred_keyp_seq.cpu().numpy() print(img_seq.shape, keypoints_seq.shape) img_seq_np = img_torch_to_numpy(img_seq) heatmaps_seq_np = heatmaps_seq.permute(0, 1, 3, 4, 2).cpu().numpy() keypoints_seq_np = keypoints_seq.cpu().numpy() d = { 'img': img_seq_np, 'keyp': keypoints_seq_np, 'heatmap': heatmaps_seq.permute(0, 1, 3, 4, 2).cpu().numpy(), 'action': data['action'].cpu().numpy() if 'action' in data else None } tmp_save_path = 'tmp_data/{}_data_{}_seed_{}'.format( l_dir, args.vids_path, args.seed) print("Save intermediate data path: ", tmp_save_path) np.savez(tmp_save_path, **d) num_seq = img_seq_np.shape[0] for i in islice(range(num_seq), 3): save_path = os.path.join( args.vids_dir, args.vids_path + "_" + l_dir + "_{}_seed_{}.mp4".format(i, args.seed)) print(i, "Video Save Path", save_path) viz_keypoints(img_seq_np[i], keypoints_seq_np[i], True, 100, save_path, args.annotate)