def discretize(F, Q, t, device): """ Discretize matrices for continuous update step with matrix exponential and matrix fraction decomposition Keyword arguments: F, Q -- KF matrices for state update (torch.Tensor) t -- time delta to last observation (torch.Tensor) device -- device of model (torch.device) Returns: A, L -- discretized update matrices """ if len(F.shape) == 3: m = F.shape[0] n = F.shape[1] M = torch.zeros(m, 2 * n, 2 * n) A = torch.matrix_exp(F * t.unsqueeze(-1).unsqueeze(-1).to(device)) M[:, :n, :n] = F M[:, :n, n:] = Q M[:, n:, n:] = -F.transpose(1, 2) M = torch.matrix_exp(M * t.unsqueeze(-1).unsqueeze(-1)) @ torch.cat( [torch.zeros(n, n), torch.eye(n, n)]) else: n = F.shape[0] M = torch.zeros(2 * n, 2 * n) A = torch.matrix_exp(F * t.view(-1, 1, 1).to(device)) M[:n, :n] = F M[:n, n:] = Q M[n:, n:] = -F.T M = torch.matrix_exp(M * t.view(-1, 1, 1)) @ torch.cat( [torch.zeros(n, n), torch.eye(n, n)]) C, D = M[:, :n], M[:, n:] L = C @ torch.inverse(D) return A.to(device), L.to(device)
def params2orb(params, coeffs, with_penalty): # params: (*, nparams) # coeffs: (*, nao, norb) nao = coeffs.shape[-2] norb = coeffs.shape[-1] nparams = params.shape[-1] bshape = params.shape[:-1] # construct the rotation parameters triu_idxs = torch.triu_indices(nao, nao, offset=1)[..., :nparams] rotmat = torch.zeros((*bshape, nao, nao), dtype=params.dtype, device=params.device) rotmat[..., triu_idxs[0], triu_idxs[1]] = params rotmat = rotmat - rotmat.transpose(-2, -1).conj() # calculate the orthogonal orbital ortho_orb = torch.matrix_exp(rotmat) @ coeffs if with_penalty: penalty = torch.zeros((1, ), dtype=params.dtype, device=params.device) return ortho_orb, penalty else: return ortho_orb
def expm(x, eps, algo='torch'): if algo == 'torch': return torch.matrix_exp(x) elif algo == 'original': return exp_from_paper(x, eps) else: raise Exception('Invalid expm algo!')
def train_exp(self, x, lie_alg_basis_ls, mat_dim, lie_var_ls, lie_alg_init_type): lie_alg_basis_ls = [p * 1. for p in lie_alg_basis_ls ] # For torch.cat, convert param to tensor. lie_alg_basis = torch.cat(lie_alg_basis_ls, dim=0)[np.newaxis, ...] # [1, lat_dim, mat_dim, mat_dim] if lie_alg_init_type == 'oth': lie_alg_basis = lie_alg_basis - lie_alg_basis.transpose(-2, -1) if self.normalize_alg: norm_v = torch.linalg.norm(lie_alg_basis, dim=[-2, -1]) lie_alg_basis = lie_alg_basis / norm_v if self.use_alg_var: lie_var = torch.cat(lie_var_ls, dim=1) # [1, lat_dim] lie_alg_basis = lie_alg_basis * lie_var[..., np.newaxis, np.newaxis] lie_group = torch.eye(mat_dim, dtype=x.dtype).to( x.device)[np.newaxis, ...] # [1, mat_dim, mat_dim] lie_alg = 0. latents_in_cut_ls = self.split_latents(x) # [x0, x1, ...] for masked_latent in latents_in_cut_ls: lie_alg_sum_tmp = torch.sum( masked_latent[..., np.newaxis, np.newaxis] * lie_alg_basis, dim=1) lie_alg += lie_alg_sum_tmp # [b, mat_dim, mat_dim] lie_group_tmp = torch.matrix_exp(lie_alg_sum_tmp) lie_group = torch.matmul(lie_group, lie_group_tmp) # [b, mat_dim, mat_dim] return lie_group
def benchmark_mexp(b, d, r, repeats=30, do_backward=True): assert d >= r if not torch.cuda.is_available(): raise RuntimeError('Benchmaring requires CUDA') eye = torch.eye(d, device='cuda:0').unsqueeze(0).repeat(b, 1, 1) start_sec = None discrepancies = [] for i in range(repeats + 1): if i == 1: torch.cuda.synchronize() start_sec = time.monotonic() param = torch.randn(b, d, d, device='cuda:0').tril(diagonal=-1) param = param - param.permute(0, 2, 1) # skew-symmetric param = torch.nn.Parameter(param) out = torch.matrix_exp(param)[:, :, :r] if do_backward: loss = out.abs().sum() loss.backward() if i >= 1: with torch.no_grad(): discrepancies.append(tensor_diff(out.permute(0, 2, 1).bmm(out), eye[:, :r, :r])) torch.cuda.synchronize() end_sec = time.monotonic() avg_ms = (end_sec - start_sec) * 1000 / repeats avg_err = torch.tensor(discrepancies).mean().item() return avg_ms, avg_err
def get_act_repr_simple(self, act_param): b, n_lats = list(act_param.size()) alg_tmp = torch.triu( torch.ones(b, n_lats, 2, 2, device=act_param.device)) alg_tmp = alg_tmp - alg_tmp.transpose(-2, -1) act_alg = act_param.view(b, n_lats, 1, 1) * alg_tmp act_repr = torch.matrix_exp(act_alg.view(-1, 2, 2)) # [b*n_lats, 2, 2] act_repr = act_repr.view(b, n_lats, 2, 2) return act_repr
def _compute_h(self): """ Compute DAG penalty. Return ------ torch.Tensor: DAG penalty term (scalar-valued). """ return torch.trace(torch.matrix_exp(self.B * self.B)) - self.d
def compute_acyclicity(w): """ Parameters ---------- w: torch.Tensor """ return torch.trace(torch.matrix_exp(w * w)) - w.shape[0]
def forward(self, source_noise, context_enc, x_prev, dx): """ Generate the output given a latent code (Inverse Normalizing Flow). Args: source_noise (FloatTensor): Source Noise (e.g., Gaussian) to do generation. context_enc (FloatTensor): Fused past_trajectory + local_scene feature. x_prev (FloatTensor): Initial positions of agents. dx (FloatTensor): Initial velocities of agents. global_scene (FloatTensor): The global scene feature map. scene_idx (IntTensor): The global_scene index corresponding to each agent. _feedback (Optional, FloatTensor): Agent states over the past T_{feedback} steps. _h (Optional, FloatTensor): GRU hidden states. Input Shapes: source_noise (A, 2) sigma (A, 2, 2) context_enc (A, D_{lc}) x_prev: (A, 2) dx: (A, 2) global_scene: (B, C, H*W) scene_idx: (A, ) _feedback (Optional): (A, 2*T_{feedback}) _h (Optional): (N_{layers}, A, D_{gru}) Output Shapes: x: (A, 2) mu: (A, 2) sigma: (A, 2, 2) """ total_agents = source_noise.size(0) # The number of agents A. output = self.mlp(context_enc) prediction = self.projection(output) # (A, 6) mu_hat = prediction[..., :2] sigma_hat = prediction[..., 2:].reshape((total_agents, 2, 2)) # Verlet integration mu = x_prev + self.velocity_const * dx + mu_hat # (A, 2) # sigma = self.symmetrize_and_exp(sigma_hat) # (A, 2, 2) sigma_sym = sigma_hat + sigma_hat.transpose(-1, -2) sigma = torch.matrix_exp( sigma_sym) # torch.matrix_exp is added in PyTorch 1.7.0 x = sigma.matmul(source_noise.unsqueeze(-1)).squeeze(-1) + mu # (A, 2) return x, mu, sigma
def blas_lapack_ops(self): m = torch.randn(3, 3) a = torch.randn(10, 3, 4) b = torch.randn(10, 4, 3) v = torch.randn(3) return ( torch.addbmm(m, a, b), torch.addmm(torch.randn(2, 3), torch.randn(2, 3), torch.randn(3, 3)), torch.addmv(torch.randn(2), torch.randn(2, 3), torch.randn(3)), torch.addr(torch.zeros(3, 3), v, v), torch.baddbmm(m, a, b), torch.bmm(a, b), torch.chain_matmul(torch.randn(3, 3), torch.randn(3, 3), torch.randn(3, 3)), # torch.cholesky(a), # deprecated torch.cholesky_inverse(torch.randn(3, 3)), torch.cholesky_solve(torch.randn(3, 3), torch.randn(3, 3)), torch.dot(v, v), torch.eig(m), torch.geqrf(a), torch.ger(v, v), torch.inner(m, m), torch.inverse(m), torch.det(m), torch.logdet(m), torch.slogdet(m), torch.lstsq(m, m), torch.lu(m), torch.lu_solve(m, *torch.lu(m)), torch.lu_unpack(*torch.lu(m)), torch.matmul(m, m), torch.matrix_power(m, 2), # torch.matrix_rank(m), torch.matrix_exp(m), torch.mm(m, m), torch.mv(m, v), # torch.orgqr(a, m), # torch.ormqr(a, m, v), torch.outer(v, v), torch.pinverse(m), # torch.qr(a), torch.solve(m, m), torch.svd(a), # torch.svd_lowrank(a), # torch.pca_lowrank(a), # torch.symeig(a), # deprecated # torch.lobpcg(a, b), # not supported torch.trapz(m, m), torch.trapezoid(m, m), torch.cumulative_trapezoid(m, m), # torch.triangular_solve(m, m), torch.vdot(v, v), )
def _expm_torch(X, basis=None): backend = utils.max_backend(X, basis) X = utils.as_tensor(X, **backend) if basis is not None: # X contains parameters in the Lie algebra -> reconstruct the matrix # X.shape = [.., F], basis.shape = [..., F, D, D] basis = utils.as_tensor(basis, **backend) X = torch.sum(basis * X[..., None, None], dim=-3) return torch.matrix_exp(X)
def forward(ctx, input): # detach so we can cast to NumPy # E = slin.expm(input.detach().numpy()) # f = np.trace(E) # E = torch.from_numpy(E) E = torch.matrix_exp(input.detach()) f = torch.trace(E) ctx.save_for_backward(E) return torch.as_tensor(f, dtype=input.dtype)
def forward(self, x, rho, alpha, temperature) -> tuple: w_prime = self._preprocess_graph(self.w, tau=temperature, seed=self.seed) w_prime = self.pns_mask * w_prime mse_loss = self._get_mse_loss(x, w_prime) h = (torch.trace(torch.matrix_exp(w_prime * w_prime)) - self.n_nodes) loss = (0.5 / self.n_samples * mse_loss + self.l1_graph_penalty * torch.linalg.norm(w_prime, ord=1) + alpha * h + 0.5 * rho * h * h) return loss, h, self.w
def infer(self, pred_traj, context_enc, x_prev, dx): """Infer the latent code given a trajectory (Normalizing Flow). Args: pred_traj (FloatTensor): Pred trajectory to do inference. encoding (FloatTensor): Context encoding. x_prev (FloatTensor): decoding position at the previous time (see the forward method). dx (FloatTensor): velocities at the previous time (see the forward method). Input Shapes: x: (A, T, 2) lc_encoding: (A, T, D_{local}) x_prev: (A, T, 2) dx: (A, T, 2) Output Shapes: z: (A, T, 2) mu: (A, T, 2) sigma: (A, T, 2, 2) """ total_agents = pred_traj.size(0) # The number of agents A. pred_steps = pred_traj.size(1) # The prediction steps T. output = self.mlp(context_enc) prediction = self.projection(output) # (A, T, 6) mu_hat = prediction[..., :2] sigma_hat = prediction[..., 2:].reshape( (total_agents, pred_steps, 2, 2)) # Verlet integration mu = x_prev + self.velocity_const * dx + mu_hat # sigma = self.symmetrize_and_exp(sigma_hat) sigma_sym = sigma_hat + sigma_hat.transpose(-1, -2) sigma = torch.matrix_exp( sigma_sym) # torch.matrix_exp is added in PyTorch 1.7.0 sigma = sigma.reshape(total_agents, pred_steps, 2, 2) # solve Z = inv(sigma) * (X-mu) x_mu = (pred_traj - mu).unsqueeze(-1) # (A, T, 2, 1) z, _ = x_mu.solve(sigma) # (A, T, 2, 1) z = z.squeeze(-1) # (A, T, 2) return z, mu, sigma
def get_act_repr(self, act_param, lie_alg_basis_ls): b, n_acts = list(act_param.size()) assert n_acts == len( lie_alg_basis_ls) # Assume each group has 1 latent dim. act_repr = torch.zeros(b, self.full_mat_dim, self.full_mat_dim).to(act_param) b_idx = 0 for i, subgroup_size_i in enumerate(self.subgroup_sizes_ls): mat_dim = int(math.sqrt(subgroup_size_i)) assert mat_dim * mat_dim == subgroup_size_i e_idx = b_idx + mat_dim assert list(lie_alg_basis_ls[i].size()) == [1, mat_dim, mat_dim] lie_alg = act_param[:, i][..., np.newaxis, np.newaxis] * ( lie_alg_basis_ls[i] - lie_alg_basis_ls[i].transpose(-2, -1) ) # Assume each latent subspace is 1, and oth basis. lie_group = torch.matrix_exp(lie_alg) # [b, mat_dim, mat_dim] act_repr[:, b_idx:e_idx, b_idx:e_idx] = lie_group b_idx = e_idx return act_repr
def sym_expm(x: torch.Tensor, using_native=False) -> torch.Tensor: r"""Symmetric matrix exponent. Parameters ---------- x : torch.Tensor symmetric matrix using_native : bool, optional if using native matrix exponent `torch.matrix_exp`, by default False Returns ------- torch.Tensor :math:`\exp(x)` """ if using_native: return torch.matrix_exp(x) else: return sym_funcm(x, torch.exp)
def forward(self, X: torch.Tensor) -> torch.Tensor: n, k = X.size(-2), X.size(-1) transposed = n < k if transposed: X = X.mT n, k = k, n # Here n > k and X is a tall matrix if self.orthogonal_map == _OrthMaps.matrix_exp or self.orthogonal_map == _OrthMaps.cayley: # We just need n x k - k(k-1)/2 parameters X = X.tril() if n != k: # Embed into a square matrix X = torch.cat( [X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1) A = X - X.mH # A is skew-symmetric (or skew-hermitian) if self.orthogonal_map == _OrthMaps.matrix_exp: Q = torch.matrix_exp(A) elif self.orthogonal_map == _OrthMaps.cayley: # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1} Id = torch.eye(n, dtype=A.dtype, device=A.device) Q = torch.linalg.solve(torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5)) # Q is now orthogonal (or unitary) of size (..., n, n) if n != k: Q = Q[..., :k] # Q is now the size of the X (albeit perhaps transposed) else: # X is real here, as we do not support householder with complex numbers A = X.tril(diagonal=-1) tau = 2. / (1. + (A * A).sum(dim=-2)) Q = torch.linalg.householder_product(A, tau) # The diagonal of X is 1's and -1's # We do not want to differentiate through this or update the diagonal of X hence the casting Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2) if hasattr(self, "base"): Q = self.base @ Q if transposed: Q = Q.mT return Q
def val_exp(self, x, lie_alg_basis_ls, lie_var_ls, lie_alg_init_type): lie_alg_basis_ls = [p * 1. for p in lie_alg_basis_ls ] # For torch.cat, convert param to tensor. lie_alg_basis = torch.cat(lie_alg_basis_ls, dim=0)[np.newaxis, ...] # [1, lat_dim, mat_dim, mat_dim] if lie_alg_init_type == 'oth': lie_alg_basis = lie_alg_basis - lie_alg_basis.transpose(-2, -1) if self.normalize_alg: norm_v = torch.linalg.norm(lie_alg_basis, dim=[-2, -1]) lie_alg_basis = lie_alg_basis / norm_v if self.use_alg_var: lie_var = torch.cat(lie_var_ls, dim=1) # [1, lat_dim] lie_alg_basis = lie_alg_basis * lie_var[..., np.newaxis, np.newaxis] lie_alg_mul = x[ ..., np.newaxis, np.newaxis] * lie_alg_basis # [b, lat_dim, mat_dim, mat_dim] lie_alg = torch.sum(lie_alg_mul, dim=1) # [b, mat_dim, mat_dim] lie_group = torch.matrix_exp(lie_alg) # [b, mat_dim, mat_dim] return lie_group
def loss_torch(p): return torch.sum((torch.matrix_exp(p) - torch.eye(3))**2)
def forward(self, X): return torch.matrix_exp(X)
def forward(self, warp_inputs: List[List[torch.Tensor]], state: List[List[torch.Tensor]] = None): embeddings = [] # Get all inputs to be same dimension so we can pass through the # transformer cascade for i, warp_input in enumerate(warp_inputs): embeddings_1 = self.input_modules[i][0](warp_input[0]) embeddings_2 = self.input_modules[i][1](warp_input[1]) embeddings += [embeddings_1, embeddings_2] embeddings = torch.cat(embeddings, dim=-1) embeddings = self.input_linear(embeddings) embeddings = self.relu(embeddings) embeddings = self.dropout(embeddings) embeddings = embeddings.transpose(0, 1) for i in range(self.transformer_layers): embeddings = self.transformer_encoder_layers[i](embeddings) embeddings = self.layer_norms[i](embeddings) embeddings = embeddings.transpose(0, 1) embeddings = embeddings.sum(dim=-2) output_raw = self.output_head(embeddings) output_raw_size = output_raw.size() output_raw = output_raw.reshape([-1, output_raw.size(-1)]) range_start = 0 kronecker_matrices = [] for i, param in enumerate(self.warp_parameters): output_matrices = [] for dim in self.output_shapes[i]: length = 2 * dim * self.num_sym_tensor_products matrix_gen = output_raw[:, range_start:range_start + length] matrix_gen = matrix_gen.reshape(-1, 2 * dim) matrix_gen_a = matrix_gen[:, :dim] matrix_gen_b = matrix_gen[:, dim:] ab = torch.bmm(matrix_gen_a.unsqueeze(-1), matrix_gen_b.unsqueeze(-2)) ba = torch.bmm(matrix_gen_b.unsqueeze(-1), matrix_gen_a.unsqueeze(-2)) matrix = 0.5 * (ab + ba) matrix = matrix.reshape(output_raw.size(0), -1, dim, dim).sum(dim=-3) matrix = matrix.reshape(list(output_raw_size)[:-1] + 2 * [dim]) exp_matrix = torch.matrix_exp( matrix.reshape((-1, matrix.size(-2), matrix.size(-1)))) exp_matrix = exp_matrix.reshape(matrix.size()) output_matrices.append(exp_matrix) range_start += length input_matrices = [] for dim in self.input_shapes[i]: length = 2 * dim * self.num_sym_tensor_products matrix_gen = output_raw[:, range_start:range_start + length] matrix_gen = matrix_gen.reshape(-1, 2 * dim) matrix_gen_a = matrix_gen[:, :dim] matrix_gen_b = matrix_gen[:, dim:] ab = torch.bmm(matrix_gen_a.unsqueeze(-1), matrix_gen_b.unsqueeze(-2)) ba = torch.bmm(matrix_gen_b.unsqueeze(-1), matrix_gen_a.unsqueeze(-2)) matrix = 0.5 * (ab + ba) matrix = matrix.reshape(output_raw.size(0), -1, dim, dim).sum(dim=-3) matrix = matrix.reshape(list(output_raw_size)[:-1] + 2 * [dim]) exp_matrix = torch.matrix_exp( matrix.reshape((-1, matrix.size(-2), matrix.size(-1)))) exp_matrix = exp_matrix.reshape(matrix.size()) input_matrices.append(exp_matrix) range_start += length kronecker_matrices.append([input_matrices, output_matrices]) return kronecker_matrices
transform=transform ) test_loader = DataLoader(test_dataset, batch_size=1, num_workers=6, shuffle=False) real_test_loader = DataLoader(real_test_dataset, batch_size=1, num_workers=6, shuffle=False) # iterating over all test and real test images, appending submission for i, data in enumerate(test_loader): filename = test_loader.dataset.sample_ids[i] img, target, K = data img = img.cuda() t, att = val_net.forward(img) # Conversion: prv --> q R_pred = torch.matrix_exp(skew_symmetric(att)) q = dcm_to_ep(R_pred) q = q.detach().cpu().numpy() if np.any(np.isnan(q)): print("Test: Nans found in q!") q = np.array([[1.0, 0.0, 0.0, 0.0]]) print(q[0]) # Conversion: [delta u, delta v, tz] -> r = [tx, ty, tz] r = origin_reg_conversion(target.detach().cpu().numpy()[0], K.detach().cpu().numpy()[0], t.detach().cpu().numpy()[0]) if np.any(np.isnan(r)): print("Test: Nans found in r!") submission.append_test(filename, q[0], r.tolist()) for i, data in enumerate(real_test_loader): filename = real_test_loader.dataset.sample_ids[i]
def compute_h(w_adj): d = w_adj.shape[0] h = torch.trace(torch.matrix_exp(w_adj * w_adj)) - d return h
while (epoch <= maxEpochs): if epoch % update_print_ct == 0: if np.abs(prev_loss - curr_loss) < eps: counteps += 1 if counteps == 3: print('The network has converged, eps = ' + str(eps)) break prev_loss = curr_loss for i in range(0, trainXp.shape[0]): dt = dt_list[i] K = torch.matrix_exp(net.Lgen * dt) # eL = torch.diag_embed(torch.exp(net.L*dt)) # exponential of the eigs, then embedded into a diagonal matrix # K = torch.matmul(torch.matmul(net.V,eL),torch.inverse(net.V)) # matrix representation of Koopman operator Kpsixp = torch.matmul(net(trainXp[i:i + 1]), K) psixf = net(trainXf[i:i + 1]) loss = loss_func(psixf, Kpsixp) optimizer.zero_grad() loss.backward() optimizer.step() curr_loss = loss.item() if epoch % update_print_ct == 0: print('[' + str(epoch) + ']' + ' loss = ' + str(curr_loss))