args.is_train = True
args.gan_mode = 'vanilla' #'vanilla'
args.trainImageSize = 128
args.skeleton_pca_dim = 52
args.skeleton_orig_dim = 63
#--device
device = torch.device("cuda:" +
args.gpu_ids if torch.cuda.is_available() else "cpu")
args.device = device
cudnn.benchmark = True
#--HPE net setting
if not 'hpe1_orig' in locals():
hpe_numBlocks = 5
hpe1_orig = dpnet(args.skeleton_pca_dim, hpe_numBlocks, device)
checkpoint = torch.load(trained_modelFile_hpe1_orig)
hpe1_orig.load_state_dict(checkpoint['state_dict'])
#--HPE2 net setting
if 'trained_modelFile_hpe2' in locals():
if not 'hpe2' in locals():
hpe2 = dpnet(args.skeleton_pca_dim, hpe_numBlocks, device)
checkpoint = torch.load(trained_modelFile_hpe2)
hpe2.load_state_dict(checkpoint['state_dict_hpe2'])
else:
if not 'hpe2' in locals():
hpe2 = dpnet(args.skeleton_pca_dim, hpe_numBlocks, device)
checkpoint = torch.load(trained_modelFile_hpe1_orig)
hpe2.load_state_dict(checkpoint['state_dict'])
args.gpu_ids = '0'
args.lambda_L1 = 50
args.trainImageSize = 128
args.skeleton_pca_dim = 52
args.skeleton_orig_dim = 63
#device
device = torch.device("cuda:%s" %
args.gpu_ids if torch.cuda.is_available() else "cpu")
args.device = device
#network
pix2pix = Pix2pix(args)
hpe1 = dpnet(52, 5, device)
hpe2 = dpnet(52, 5, device)
net = FusionNet(pix2pix, hpe1, hpe2, args)
w = np.random.rand(args.skeleton_orig_dim, args.skeleton_pca_dim)
b = np.random.rand(args.skeleton_orig_dim)
net.set_reconstruction_net(w, b)
net = net.to(device)
#run
ir = np.random.rand(args.train_batch, 1, 128, 128)
depth = np.random.rand(args.train_batch, 1, 128, 128)
net.set_input(ir, depth)