Spaces:
Runtime error
Runtime error
| from torch import distributions | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader | |
| from torch.distributions import Normal | |
| import numpy as np | |
| import cv2 | |
| import trimesh | |
| from tqdm import tqdm | |
| import warnings | |
| warnings.filterwarnings("ignore", category=DeprecationWarning) | |
| import FrEIA.framework as Ff | |
| import FrEIA.modules as Fm | |
| from configs.barc_cfg_defaults import get_cfg_global_updated | |
| class NormalizingFlowPrior(nn.Module): | |
| def __init__(self, nf_version=None): | |
| super(NormalizingFlowPrior, self).__init__() | |
| # the normalizing flow network takes as input a vector of size (35-1)*6 which is | |
| # [all joints except root joint]*6. At the moment the rotation is represented as 6D | |
| # representation, which is actually not ideal. Nevertheless, in practice the | |
| # results seem to be ok. | |
| n_dim = (35 - 1) * 6 | |
| self.param_dict = self.get_version_param_dict(nf_version) | |
| self.model_inn = self.build_inn_network(n_dim, k_tot=self.param_dict['k_tot']) | |
| self.initialize_with_pretrained_weights() | |
| def get_version_param_dict(self, nf_version): | |
| # we had trained several version of the normalizing flow pose prior, here we just provide | |
| # the option that was user for the cvpr 2022 paper (nf_version=3) | |
| if nf_version == 3: | |
| param_dict = { | |
| 'k_tot': 2, | |
| 'path_pretrained': get_cfg_global_updated().paths.MODELPATH_NORMFLOW, | |
| 'subnet_fc_type': '3_64'} | |
| else: | |
| print(nf_version) | |
| raise ValueError | |
| return param_dict | |
| def initialize_with_pretrained_weights(self, weight_path=None): | |
| # The normalizing flow pose prior is pretrained separately. Afterwards all weights | |
| # are kept fixed. Here we load those pretrained weights. | |
| if weight_path is None: | |
| weight_path = self.param_dict['path_pretrained'] | |
| print(' normalizing flow pose prior: loading {}..'.format(weight_path)) | |
| pretrained_dict = torch.load(weight_path)['model_state_dict'] | |
| self.model_inn.load_state_dict(pretrained_dict, strict=True) | |
| def subnet_fc(self, c_in, c_out): | |
| if self.param_dict['subnet_fc_type']=='3_512': | |
| subnet = nn.Sequential(nn.Linear(c_in, 512), nn.ReLU(), | |
| nn.Linear(512, 512), nn.ReLU(), | |
| nn.Linear(512, c_out)) | |
| elif self.param_dict['subnet_fc_type']=='3_64': | |
| subnet = nn.Sequential(nn.Linear(c_in, 64), nn.ReLU(), | |
| nn.Linear(64, 64), nn.ReLU(), | |
| nn.Linear(64, c_out)) | |
| return subnet | |
| def build_inn_network(self, n_input, k_tot=12, verbose=False): | |
| coupling_block = Fm.RNVPCouplingBlock | |
| nodes = [Ff.InputNode(n_input, name='input')] | |
| for k in range(k_tot): | |
| nodes.append(Ff.Node(nodes[-1], | |
| coupling_block, | |
| {'subnet_constructor':self.subnet_fc, 'clamp':2.0}, | |
| name=F'coupling_{k}')) | |
| nodes.append(Ff.Node(nodes[-1], | |
| Fm.PermuteRandom, | |
| {'seed':k}, | |
| name=F'permute_{k}')) | |
| nodes.append(Ff.OutputNode(nodes[-1], name='output')) | |
| model = Ff.ReversibleGraphNet(nodes, verbose=verbose) | |
| return model | |
| def calculate_loss_from_z(self, z, type='square'): | |
| assert type in ['square', 'neg_log_prob'] | |
| if type == 'square': | |
| loss = (z**2).mean() # * 0.00001 | |
| elif type == 'neg_log_prob': | |
| means = torch.zeros((z.shape[0], z.shape[1]), dtype=z.dtype, device=z.device) | |
| stds = torch.ones((z.shape[0], z.shape[1]), dtype=z.dtype, device=z.device) | |
| normal_distribution = Normal(means, stds) | |
| log_prob = normal_distribution.log_prob(z) | |
| loss = - log_prob.mean() | |
| return loss | |
| def calculate_loss(self, poses_rot6d, type='square'): | |
| assert type in ['square', 'neg_log_prob'] | |
| poses_rot6d_noglob = poses_rot6d[:, 1:, :].reshape((-1, 34*6)) | |
| z, _ = self.model_inn(poses_rot6d_noglob, rev=False, jac=False) | |
| loss = self.calculate_loss_from_z(z, type=type) | |
| return loss | |
| def forward(self, poses_rot6d): | |
| # from pose to latent pose representation z | |
| # poses_rot6d has shape (bs, 34, 6) | |
| poses_rot6d_noglob = poses_rot6d[:, 1:, :].reshape((-1, 34*6)) | |
| z, _ = self.model_inn(poses_rot6d_noglob, rev=False, jac=False) | |
| return z | |
| def run_backwards(self, z): | |
| # from latent pose representation z to pose | |
| poses_rot6d_noglob, _ = self.model_inn(z, rev=True, jac=False) | |
| return poses_rot6d_noglob | |