def extract_features(self, x): coords = x['coordinates'] sinput, inds_list = self.create_sparse_tensors(coords) features_dict = self.feature_extractor(sinput) features = features_dict["features"] if "attention" in features_dict.keys(): att = features_dict["attention"] batch_indices = list(features.coords_man.get_batch_indices()) features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: mask = features.C[:, 0] == batch_indices[ind_cnt] # hacky way of finding batch channel dimension if mask.int().sum() != inds_list[ind_cnt].shape[0]: mask = features.C[:, -1] == batch_indices[ind_cnt] f = features.F[mask] assert f.shape[0] == inds_list[ind_cnt].shape[0] if "attention" in features_dict.keys(): a = att.F[mask] assert a.shape[0] == inds_list[ind_cnt].shape[0] att_L.append(a) features_L.append(f.permute(1, 0)) coords_ds_L.append(coords_s[:, inds_list[ind_cnt]]) inds_L.append(inds_list[ind_cnt]) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) out = dict() out['features'] = features_LL out['att'] = att_LL out['coordinates_ds'] = coords_ds_LL.to(self.params.device) out['indices'] = inds_LL out['coordinates'] = x['coordinates'] return out
def forward(self, x): coords = x['coordinates'].clone() t_tot = time.time() features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: pcd_down, f = self.preprocess_point_cloud( coords_s, self.params.voxel_size) features_L.append(f) coords_ds_L.append(pcd_down) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) x = dict() x['features'] = self.cluster_features( features_LL, self.params.feature_distr_parameters.num_feature_clusters) x['att'] = att_LL x['coordinates_ds'] = coords_ds_LL.to(self.params.device) x['indices'] = inds_LL out = self.registration(x) tot_time = time.time() - t_tot print("tot time: %.1f ms" % (tot_time * 1000)) out["time"] = tot_time out["features"] = features_LL out["indices"] = inds_LL out["coordinates_ds"] = coords_ds_LL return out
class MultinomialModel(BaseFeatureModel): def __init__(self, parameters, K, features, repeat_list): self.params = parameters self.repeat_list = repeat_list self.initialize(K, features) def initialize(self, K, features): init_method = self.params.get("init_method", "dirichlet") if init_method == "uniform": init_distr_list = [] for i in range(len(features)): init_distr_list.append( torch.ones(K, self.params.num_feature_clusters).float() / self.params.num_feature_clusters) elif init_method == "dirichlet": init_distr_list = TensorList() relative_distr = features.sum(dim=1).sum_list() relative_distr = relative_distr / relative_distr.sum() for r in relative_distr: dir = torch.distributions.dirichlet.Dirichlet(r) init_distr_list.append(dir.sample((K, ))) else: init_distr_list = [] for i in range(len(features)): init_distr_list.append( torch.ones(K, self.params.num_feature_clusters).float() / self.params.num_feature_clusters) self.distr = TensorListList(init_distr_list, repeat=self.repeat_list) def to(self, dev): self.distr = self.distr.to(dev) def posteriors(self, y): p = y.permute(1, 0) @ self.distr.permute(1, 0) # marginalization return p def maximize(self, ap, ow, y, den): tmp = y / den.permute(1, 0) as_sum = (tmp @ (ow * ap)).sum_list() self.distr = self.distr * TensorListList(as_sum.permute(1, 0), repeat=self.repeat_list) self.distr = self.distr / self.distr.sum(dim=1, keepdims=True)
class VonMisesModelList(BaseFeatureModel): def __init__(self, parameters, K, features, s, mu, repeat_list): self.params = parameters self.mu = self.initialize_mu(K, features) if len(mu) > 0: self.mu = TensorListList(mu, repeat=repeat_list) self.K = K self.repeat_list = repeat_list self.s2 = s * s self.local_posterior = 1 def to(self, dev): self.mu = self.mu.to(dev) def initialize_mu(self, K, features): X = TensorList() for TV in features: Xi = np.random.randn(TV[0].shape[0], K).astype(np.float32) Xi = torch.from_numpy(Xi).to(TV[0].device) Xi = Xi / torch.norm(Xi, dim=0, keepdim=True) X.append(Xi.permute(1, 0)) return X def posteriors(self, y): log_p = y.permute(1, 0) @ TensorListList(self.mu.permute(1, 0), repeat=self.repeat_list) p = log_p / self.s2 return p.exp() def maximize(self, a, y, den): self.mu = ((y @ a).sum_list()).permute(1, 0) self.mu = self.mu / self.mu.norm(dim=-1, keepdim=True) return def detach(self): self.mu = self.mu.detach()
def register_point_sets(self, x): Vs = x["coordinates_ds"] features = x["features"] features_w = x["att"] repeat_list = [len(V) for V in Vs] init_R, init_t = get_init_transformation_list( Vs, self.params.get("mean_init", True)) TVs = init_R @ Vs + init_t X = TensorList() Q = TensorList() mu = TensorList() for TV, Fs in zip(TVs, features): if self.params.cluster_init == "box": Xi = get_randn_box_cluster_means_list(TV, self.params.K) else: Xi = get_randn_sphere_cluster_means_list( TV, self.params.K, self.params.get("cluster_mean_scale", 1.0)) Q.append( get_scaled_cluster_precisions_list( TV, Xi, self.params.cluster_precision_scale)) X.append(Xi.T) feature_distr = feature_models.MultinomialModel( self.params.feature_distr_parameters, self.params.K, features, repeat_list=repeat_list) feature_distr.to(self.params.device) X = TensorListList(X, repeat=repeat_list) self.betas = get_default_beta(Q, self.params.gamma) Vs = Vs TVs = TVs # Compute the observation weights if self.params.use_dare_weighting: observation_weights = empirical_estimate(Vs, self.params.ow_args) ow_reg_factor = 8.0 ow_mean = observation_weights.mean(dim=0, keepdim=True) for idx in range(len(observation_weights)): for idxx in range(len(observation_weights[idx])): observation_weights[idx][idxx][observation_weights[idx][idxx] > ow_reg_factor * ow_mean[idx][idxx]] \ = ow_reg_factor * ow_mean[idx][idxx] else: observation_weights = 1.0 ds = TVs.permute(1, 0).sqe(X).permute(1, 0) if self.params.debug: self.visdom.register( dict(pcds=Vs[0].cpu(), X=X[0][0].cpu(), c=None), 'point_clouds', 2, 'init') time.sleep(1) Rs = init_R.to(self.params.device) ts = init_t.to(self.params.device) self.betas = TensorListList(self.betas, repeat=repeat_list) QL = TensorListList(Q, repeat=repeat_list) Riter = TensorListList() titer = TensorListList() TVs_iter = TensorListList() for i in range(self.params.num_iters): Qt = QL.permute(1, 0) ap = (-0.5 * ds * QL).exp() * QL.pow(1.5) if i < 1000: pyz_feature = feature_distr.posteriors(features) else: pyz_feature = 1.0 a = ap * pyz_feature ac_den = a.sum(dim=-1, keepdim=True) + self.betas a = a / ac_den # normalize row-wise a = a * observation_weights L = a.sum(dim=-2, keepdim=True).permute(1, 0) W = (Vs @ a) * QL b = L * Qt # weights, b mW = W.sum(dim=-1, keepdim=True) mX = (b.permute(1, 0) @ X).permute(1, 0) z = L.permute(1, 0) @ Qt P = (W @ X).permute(1, 0) - mX @ mW.permute(1, 0) / z # Compute R and t svd_list_list = P.cpu().svd() Rs = TensorListList() for svd_list in svd_list_list: Rs_list = TensorList() for svd in svd_list: uu, vv = svd.U, svd.V vvt = vv.permute(1, 0) detuvt = uu @ vvt detuvt = detuvt.det() S = torch.ones(1, 3) S[:, -1] = detuvt Rs_list.append((uu * S) @ vvt) Rs.append(Rs_list) Rs = Rs.to(self.params.device) Riter.append(Rs) ts = (mX - Rs @ mW) / z titer.append(ts) TVs = Rs @ Vs + ts TVs_iter.append(TVs.clone()) if self.params.debug: self.visdom.register( dict(pcds=TVs[0].cpu(), X=X[0][0].cpu(), c=None), 'point_clouds', 2, 'registration-iter') time.sleep(0.2) # Update X den = L.sum_list() if self.params.fix_cluster_pos_iter < i: X = (TVs @ a).permute(1, 0) X = TensorListList(X.sum_list() / den, repeat_list) # Update Q ds = TVs.permute(1, 0).sqe(X).permute(1, 0) wn = (a * ds).sum(dim=-2, keepdim=True).sum_list() Q = (3 * den / (wn.permute(1, 0) + 3 * den * self.params.epsilon)).permute( 1, 0) QL = TensorListList(Q, repeat=repeat_list) feature_distr.maximize(ap=ap, ow=observation_weights, y=features, den=ac_den) out = dict(Rs=Rs, ts=ts, X=X, Riter=Riter[:-1], titer=titer[:-1], Vs=TVs, Vs_iter=TVs_iter[:-1], ow=observation_weights) return out
def forward(self, x): coords = x['coordinates'].clone() t_tot = time.time() sinput, inds_list = self.create_sparse_tensors(coords) features_dict = self.feature_extractor(sinput) features = features_dict["features"] if "attention" in features_dict.keys(): att = features_dict["attention"] if torch.isnan(features.feats).any(): print("nans in features!") batch_indices = list(features.coords_man.get_batch_indices()) features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: mask = features.C[:, 0] == batch_indices[ind_cnt] if mask.int().sum() != inds_list[ind_cnt].shape[0]: mask = features.C[:, -1] == batch_indices[ind_cnt] f = features.F[mask] assert f.shape[0] == inds_list[ind_cnt].shape[0] if "attention" in features_dict.keys(): a = att.F[mask] assert a.shape[0] == inds_list[ind_cnt].shape[0] att_L.append(a) features_L.append(f) coords_ds_L.append(coords_s[:, inds_list[ind_cnt]]) inds_L.append(inds_list[ind_cnt]) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) x = dict() x['features'] = self.cluster_features( features_LL, self.params.feature_distr_parameters.num_feature_clusters) x['att'] = att_LL x['coordinates_ds'] = coords_ds_LL.to(self.params.device) x['indices'] = inds_LL out = self.registration(x) tot_time = time.time() - t_tot print("tot time: %.1f ms" % (tot_time * 1000)) out["time"] = tot_time out["features"] = features_LL out["indices"] = inds_LL out["coordinates_ds"] = coords_ds_LL return out
def forward(self, x_in): coords = x_in['coordinates'].clone() t_tot = time.time() sinput, inds_list = self.create_sparse_tensors(coords) resample_time = time.time() - t_tot print("resample time: %.1f ms" % ((resample_time) * 1000)) if not self.params.feature_distr_parameters.model=='none': features_dict = self.feature_extractor(sinput) features = features_dict["features"] if "attention" in features_dict.keys(): att = features_dict["attention"] extract_time=time.time()-t_tot print("extract time: %.1f ms" % ((extract_time) * 1000)) batch_indices = list(features.coords_man.get_batch_indices()) else: features_dict=None extract_time=0 time_preprocess = time.time() features_LL = TensorListList() att_LL = TensorListList() coords_ds_LL = TensorListList() inds_LL = [] ind_cnt = 0 for coords_b in coords: features_L = TensorList() att_L = TensorList() coords_ds_L = TensorList() inds_L = [] for coords_s in coords_b: if not features_dict is None: mask = features.C[:, 0] == batch_indices[ind_cnt] # hacky way of finding batch channel dimension if mask.int().sum() != inds_list[ind_cnt].shape[0]: mask = features.C[:, -1] == batch_indices[ind_cnt] f = features.F[mask] assert f.shape[0] == inds_list[ind_cnt].shape[0] if "attention" in features_dict.keys(): a = att.F[mask] assert a.shape[0] == inds_list[ind_cnt].shape[0] att_L.append(a) features_L.append(f.permute(1, 0)) coords_ds_L.append(coords_s[:, inds_list[ind_cnt]]) inds_L.append(inds_list[ind_cnt]) ind_cnt = ind_cnt + 1 features_LL.append(features_L) att_LL.append(att_L) coords_ds_LL.append(coords_ds_L) inds_LL.append(inds_L) x = dict() x['features'] = features_LL x['att'] = att_LL x['coordinates_ds'] = coords_ds_LL.to(self.params.device) x['indices'] = inds_LL x['coordinates'] = x_in['coordinates'] print("preprocess time: %.1f ms" % ((time.time()-time_preprocess) * 1000)) reg_time=time.time() out = self.registration(x) reg_time2 = time.time() - reg_time print("reg time: %.1f ms" % ((time.time() - reg_time) * 1000)) tot_time = time.time() - t_tot print("tot time: %.1f ms" % (tot_time * 1000)) out["time"] = tot_time out["reg_time"] = reg_time2 out["extract_time"] =extract_time out["resample_time"] = resample_time out["coordinates_ds"] = coords_ds_LL return out
def __call__(self, batch, epoch): if self.epoch != epoch: self.epoch = epoch self.iter_cnt = 0 self.num_success_acc = 0 data, info = batch batch_size = len(data['coordinates']) device = self.model.params.device loss_iter = 0 loss = 0 final_loss = 0 out_feat = self.model.extract_features(data) features = out_feat["features"] info["R_gt"] = info["R_gt"].to(device) info["t_gt"] = info["t_gt"].to(device) num_pairs = 0 for fb in features: M = len(fb) for ind1 in range(M - 1): for ind2 in range(ind1 + 1, M): num_pairs = num_pairs + 1 # check valid pairs wrt number of correspondences # invalid pairs are ignored in the computation of the loss valid_pairs_LL = [] if self.min_corresponence_rate > 0.0: correspondences = [] for coord_ds, r, t in zip(out_feat["coordinates_ds"], info["R_gt"], info["t_gt"]): correspondences.append( utils.extract_correspondences_gpu(coord_ds, r, t)) for fb, corr in zip(features, correspondences): db = corr["distances"].to(device) M = len(fb) cnt = 0 valid_pairs_L = [] for ind1 in range(M - 1): for ind2 in range(ind1 + 1, M): mask = db[cnt] < self.th corrs_rate = mask.sum().float() / mask.shape[0] valid_pairs_L.append( corrs_rate > self.min_corresponence_rate) cnt = cnt + 1 valid_pairs_LL.append(sum(valid_pairs_L)) else: for fb in features: cnt = 0 M = len(fb) for ind1 in range(M - 1): for ind2 in range(ind1 + 1, M): cnt = cnt + 1 valid_pairs_LL.append(cnt) Rgt, tgt = info["R_gt"], info["t_gt"] num_samples = sum(v > 0.0 for v in valid_pairs_LL) if self.min_corresponence_rate > 0.0: self.iter_cnt += num_samples.item() else: self.iter_cnt += num_samples # only compute loss is there is at least one valid pair in the batch if sum(valid_pairs_LL) == 0: return loss, self.num_success_acc / (self.iter_cnt) if sum(valid_pairs_LL) < num_pairs: out_feat_filt = dict() for k in out_feat.keys(): out_feat_filt[k] = TensorListList([ out_feat[k][i] for i in range(len(out_feat[k])) if valid_pairs_LL[i] ]) Rgt = TensorListList( [Rgt[i] for i in range(batch_size) if valid_pairs_LL[i]]) tgt = TensorListList( [tgt[i] for i in range(batch_size) if valid_pairs_LL[i]]) out_feat = out_feat_filt out_reg = self.model.register(out_feat) if not self.vis is None: self.vis(out_reg, info, data) Rs, ts = out_reg["Rs"], out_reg["ts"] Riter, titer = out_reg["Riter"], out_reg["titer"] Rgt = Rgt.to(Rs[0][0].device) tgt = tgt.to(ts[0][0].device) rot_errs = utils.L2sqrList(Rs.detach(), Rgt).sqrt() trans_errs = utils.L2sqrTransList(ts.detach(), tgt, Rs.detach(), Rgt).sqrt() # check number of successful registrations num_success = 0 valid_list = [] for rs_err, ts_err in zip(rot_errs, trans_errs): for rerr, terr in zip(rs_err, ts_err): val = (rerr < self.eval_rot_err_thresh) * ( terr < self.eval_trans_err_thresh) valid_list.append(val.item()) num_success = num_success + val.item() self.num_success_acc += num_success # compute registration error per iteration for Rit, tit, w in zip( Riter, titer, self.weight["Vs_iter"][self.compute_loss_iter:]): if w > 0: trans_errs2 = utils.GARPointList(tit, tgt, Rit, Rgt, V=out_feat["coordinates_ds"], c=self.c, alpha=self.alpha) for terr in trans_errs2: for t in terr: loss_iter = loss_iter + w * t # compute final registration error if self.weight["Vs"] > 0.0: trans_errs2 = utils.GARPointList(ts, tgt, Rs, Rgt, V=out_feat["coordinates_ds"], c=self.c, alpha=self.alpha) for terr in trans_errs2: for t in terr: final_loss = final_loss + self.weight["Vs"] * t print("num valid: ", num_success, "num_success_acc rate: ", self.num_success_acc / (self.iter_cnt)) if sum(valid_pairs_LL): loss = loss + final_loss + loss_iter loss = loss / num_samples return loss, self.num_success_acc / (self.iter_cnt)