import torch import torch.nn as nn from copy import deepcopy import torch.nn.functional as F from einops import rearrange from torch_cluster import fps from os import sys, path import math import numpy as np from collections import (defaultdict, OrderedDict, deque) import os from SoftGroup.softgroup.ops import (voxelization_idx, voxelization) from model.basic_vae import (Encoder, Decoder) from model.NGC import get_NGC_structure from model.common import (ResidualConv, CappedLayerNorm) import pyvista as pv import matplotlib.pyplot as plt # Chamfer distance calculation class VQVAE(nn.Module): def __init__(self, input_dim, hidden_dim, codebook_size, embedding_dim, num_points, voxel_size, spconv_channels, blocks_num, smooth_end_epoch, beta, lambda_chamfer, lambda_vq_start, lambda_vq_final, lambda_construction, lambda_usage, warmup_steps, ema_decay, token_N1_stage, token_N2_stage, token_N3_stage, decoder_layer, buffer_stage_num): super(VQVAE, self).__init__() self.device = device = torch.device("cuda") self.blocks_num = blocks_num self.spconv_channels = spconv_channels self.encoder = Encoder(input_dim, hidden_dim, self.spconv_channels, self.blocks_num) self.hidden_dim = hidden_dim self.embedding_dim = embedding_dim self.codebook_size = codebook_size self.num_points = num_points self.voxel_size = voxel_size self.side_length_patchs = [(2 ** i) for i in range(self.voxel_size)] self.token_N3_stage = token_N3_stage self.token_N2_stage = token_N2_stage self.token_N1_stage = token_N1_stage self.token_N1_nums = 2 ** (token_N1_stage - 1) self.token_N2_nums = 2 ** (token_N2_stage - 1) self.token_N3_nums = 2 ** (token_N3_stage - 1) self.token_stage = self.token_N2_stage + buffer_stage_num self.token_stage_nums_list = [(2 ** i) for i in range(self.token_stage)] self.if_normal_compress = True self.num_freqs = 6 self.pos_dim = 3 + 2 * 3 * self.num_freqs self.pos_linear = nn.Linear(self.pos_dim, self.spconv_channels) self.after_spconv = nn.Sequential( nn.Linear(spconv_channels, hidden_dim), nn.LayerNorm(hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, embedding_dim), CappedLayerNorm(embedding_dim) ) nn.init.xavier_uniform_(self.pos_linear.weight) nn.init.zeros_(self.pos_linear.bias) self.embeddings = nn.Embedding(codebook_size, embedding_dim) self.embeddings.weight.data.uniform_(-1.0 / embedding_dim , 1.0 / embedding_dim) self.smooth_end_epoch = smooth_end_epoch self.phi = nn.ModuleList([ ResidualConv(embedding_dim, embedding_dim) for _ in range(self.token_stage) ]) self.beta = beta self.decoder = Decoder(embedding_dim, hidden_dim, input_dim, decoder_layer) # 6 layers self.if_freeze_spconv = False self.lambda_chamfer = lambda_chamfer self.warmup_steps = warmup_steps self.lambda_vq_start = lambda_vq_start self.lambda_vq_final = lambda_vq_final self.store_step = 200 self.epoch_step = 20 self.lambda_construction = lambda_construction self.alpha = nn.Parameter(torch.tensor(0.1)) self.act = nn.GELU() self.log_data = torch.zeros(8, dtype=torch.float32) self.usage_weight = lambda_usage self.ema_decay = ema_decay self.vq_window = deque(maxlen=10) # window_size e.g. 10 self.triggers = 0 self.patience_counter = 0 self.safe_vq_pct_thr = 0.40 self.lowest_vq_weight = 0.03 self.init_ema() def init_ema(self): # Ensure embeddings won't be touched by optimizer self.embeddings.weight.requires_grad = False # register or create EMA buffers if not exist # Use register_buffer if inside nn.Module; otherwise set attributes if not hasattr(self, 'cluster_size'): self.register_buffer('cluster_size', torch.zeros(self.codebook_size, dtype=torch.float32)) if not hasattr(self, 'embed_avg'): self.register_buffer('embed_avg', torch.zeros(self.codebook_size, self.embedding_dim, dtype=torch.float32)) # move to device if user passed one self.cluster_size = self.cluster_size.to(self.device) self.embed_avg = self.embed_avg.to(self.device) self.ema_cluster_size = None self.ema_embed_sum = None self.ema_weights_list = [((i + 1) / self.token_stage) for i in range(self.token_stage)] def init_from_ckpt(self, path, ignore_keys, if_only_init_param): ignore_keys=list() sd = torch.load(path, map_location=self.device) if not if_only_init_param: # return opt & curr_epoch curr_epoch = sd['epoch'] params = sd['model'] opt = sd['optimizer'] missing, unexpected = self.load_state_dict(params, strict=False) print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") if len(missing) > 0: print(f"Missing keys: {missing}") if len(unexpected) > 0: print(f"Unexpected keys: {unexpected}") return opt, curr_epoch else: # start with 0 epoch params = sd['model'] keys = list(params.keys()) missing, unexpected = self.load_state_dict(params, strict=False) print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") if len(missing) > 0: print(f"Missing keys: {missing}") if len(unexpected) > 0: print(f"Unexpected keys: {unexpected}") return None, None @torch.no_grad() def update_ema(self, decay, eps=1e-5): """ z_e: (B, N, D) encoder outputs (float) codes: (B, N) long tensor of indices in [0, num_embeddings) This updates self.cluster_size and self.embed_avg, then writes to embeddings. """ cluster_size_batch = self.ema_cluster_size embed_sum_batch = self.ema_embed_sum # EMA update (in-place) self.cluster_size.mul_(decay).add_(cluster_size_batch, alpha=1.0 - decay) # (K,) self.embed_avg.mul_(decay).add_(embed_sum_batch, alpha=1.0 - decay) # (K, D) # compute normalized cluster sizes to avoid divide-by-zero n = self.cluster_size.sum() # stabilized cluster_size for division cluster_size_normalized = ((self.cluster_size + eps) / (n + self.codebook_size * eps)) * n # (K,) # new embeddings = embed_avg / cluster_size_normalized.unsqueeze(1) # guard against zeros/nans denom = cluster_size_normalized.unsqueeze(1) # (K,1) # safe division new_emb = self.embed_avg / denom # replace any NaN/Inf entries with current embedding values (safe fallback) if torch.isnan(new_emb).any() or torch.isinf(new_emb).any(): cur = self.embeddings.weight.data.to(device) nan_mask = (~torch.isfinite(new_emb)) new_emb[nan_mask] = cur[nan_mask] # write back to embedding weights (in-place) self.embeddings.weight.data.copy_(new_emb) def point_cloud_normalize(self, x): x_normalized = None if self.if_normal_compress: x_max_minus_min = x.max(dim=1, keepdim=True)[0] - x.min(dim=1, keepdim=True)[0] + 1e-8 max_val = x_max_minus_min.max(dim=-1, keepdim=True)[0] x_norm = x_max_minus_min / (max_val + 1e-8) x_normalized = ((x - x.min(dim=1, keepdim=True)[0]) / x_max_minus_min) * x_norm pass else: x_normalized = (x - x.min(dim=1, keepdim=True)[0]) / (x.max(dim=1, keepdim=True)[0] - x.min(dim=1, keepdim=True)[0] + 1e-8) # [0, 1] scaled_x_normalized = x_normalized * (self.side_length_patchs[-1] - 1) # [0, 63] x_normalized = 2. * x_normalized - 1. # [-1, 1] return x_normalized, scaled_x_normalized def point_cloud_to_voxel(self, x): # smn B, N, C = x.shape # Normalize point cloud to [0, resolution-1] x_normalized, scaled_x_normalized = self.point_cloud_normalize(x) # x_normalized [-1, 1], scaled_x_normalized [0, resolution-1] x_normalized_long = scaled_x_normalized.long().cpu() batch_ids = torch.arange(B, dtype=torch.long).view(B, 1, 1).expand(B, N, 1) coords_long = torch.cat([batch_ids, x_normalized_long], dim=-1) coords_long = rearrange(coords_long, 'b n c -> (b n) c') batch_ids = batch_ids.to(x_normalized.device) coords_float = torch.cat([batch_ids, scaled_x_normalized], dim=-1) coords_float = rearrange(coords_float, 'b n c -> (b n) c') voxel_coords, v2p_map, p2v_map = voxelization_idx(coords_long, B) spatial_shape = [self.side_length_patchs[-1]] * 3 p2v_map = p2v_map.to(x_normalized.device) voxel_feats = voxelization(coords_float, p2v_map) # mode=4 avg batch_ids = voxel_feats[:, 0].long() voxel_feats = voxel_feats[:, 1:] voxel_feats_split = [voxel_feats[batch_ids == b] for b in range(B)] voxel_coords_split = [] for i in range(B): voxel_coords_split.append(voxel_feats_split[i].long()) return voxel_feats_split, voxel_coords_split, spatial_shape, x_normalized def fourier_embed(self, x, num_freqs): x_norm = self.normalize_to_minus1_1(x) freqs = 2 ** torch.arange(num_freqs, device=x_norm.device) * torch.pi x_proj = x_norm[..., None] * freqs embed = torch.cat([x_norm, torch.sin(x_proj).flatten(-2), torch.cos(x_proj).flatten(-2)], dim=-1) embed = embed * (self.side_length_patchs[-1] / 2.0) return embed def flat_tokens(self, voxel_feats_total, voxel_batch_id_total, voxel_floatcoords_total): # first FPS to 8192 token nums L = self.token_N3_nums B = len(voxel_feats_total) sampled_feats_total = [] sampled_floatcoords_total = [] for batch_id in range(B): voxel_batch_id = voxel_batch_id_total[batch_id].squeeze(1).long() voxel_floatcoords = voxel_floatcoords_total[batch_id] voxel_feats = voxel_feats_total[batch_id].features points_num = voxel_batch_id.shape[0] if points_num < L: sample_idx = np.random.choice(points_num, L, replace=True) else: sample_idx = fps(voxel_floatcoords, voxel_batch_id, ratio=L/points_num) sampled_feats = voxel_feats[sample_idx] sampled_floatcoords = voxel_floatcoords[sample_idx] sampled_feats_total.append(sampled_feats) sampled_floatcoords_total.append(sampled_floatcoords) sampled_feats_total = torch.stack(sampled_feats_total, dim=0) sampled_floatcoords_total = torch.stack(sampled_floatcoords_total, dim=0) floatcoords_total = sampled_floatcoords_total # pos embbedding & after-encode token mlp pos_embedding = self.act(self.pos_linear(self.fourier_embed(sampled_floatcoords_total, self.num_freqs))) token_feats = sampled_feats_total + self.alpha * pos_embedding # [-1 1] token_feats = self.after_spconv(token_feats) # prepare for the construction loss norm_floatcoords_total = self.normalize_to_minus1_1(floatcoords_total) return token_feats, norm_floatcoords_total def encode(self, voxel_feats, voxel_coords, spatial_shape, epoch): # first voxelize & spconv encode pc voxel_feats_total, voxel_coords_total, voxel_batch_id_total = self.encoder(voxel_feats, voxel_coords, spatial_shape) # [0, resolution-1] # next flat voxel_feats into N3 tokens token_feats_N3, floatcoords_N3 = self.flat_tokens(voxel_feats_total, voxel_batch_id_total, voxel_feats) # [-1 1] # get structure gt & N3 -> N2 structure_gt02_bin, structure_gt02_dec, token_feats_N2, floatcoords_N2 = get_NGC_structure(token_feats_N3, floatcoords_N3, self.token_N2_nums, self.token_N1_nums, self.token_stage) B, L, C = token_feats_N2.shape for b in range(B): structure_gt02_dec[b] = structure_gt02_dec[b][:, :self.token_stage] structure_gt02_dec = torch.stack(structure_gt02_dec, dim=0).permute(0, 2, 1) # codebook_size = K, embedding_dim = D self.ema_cluster_size = torch.zeros(self.codebook_size, device=self.device) self.ema_embed_sum = torch.zeros(self.codebook_size, C, device=self.device) # start quantizing f_BLC = token_feats_N2 f_no_grad = f_BLC.detach() f_rest = f_no_grad.clone() f_hat = torch.zeros_like(f_rest) embedding = self.embeddings.weight min_encoding_stages_indices = [] SN = len(self.token_stage_nums_list) mean_vq_loss: torch.Tensor = 0.0 vq_loss_dict = defaultdict(float) for si in range(SN): # find the nearest embedding if si == SN - 1: # last stage structure_map = torch.arange(L, device=self.device, dtype=torch.long).unsqueeze(0).repeat(B, 1) else: structure_map = structure_gt02_dec[:, si, :] # [B, L] h_BLC_list = [] min_encoding_indices = [] for b in range(B): structure_b = structure_map[b] # (L,) f_rest_b = f_rest[b] # (L, C) uniq, inv = torch.unique(structure_b, return_inverse=True) M = uniq.numel() # unique class_sums = torch.zeros(M, C, device=f_rest.device) class_sums.scatter_add_(0, inv.unsqueeze(-1).expand(-1, C), f_rest_b) # (M, C) counts = torch.zeros(M, device=f_rest.device) counts.scatter_add_(0, inv, torch.ones_like(inv, dtype=counts.dtype)) # (M,) rest_NC_b = class_sums / counts.unsqueeze(-1) # (M, C) d = ( rest_NC_b.pow(2).sum(1, keepdim=True) + embedding.pow(2).sum(1) - 2 * rest_NC_b @ embedding.T ) # (M, K) if epoch is not None and self.smooth_end_epoch != -1 and epoch < self.smooth_end_epoch: logit = F.softmax(d.max(dim=-1, keepdim=True)[0] - d, dim=-1) idx = torch.argmax(logit, dim=-1) one_hot = F.one_hot(idx, self.codebook_size).type_as(logit) one_hot = one_hot - logit.detach() + logit h_NC_b = one_hot @ embedding else: idx = torch.argmin(d, dim=1) # M h_NC_b = embedding[idx] # store some ema information # rest_NC: (Mi, C) # idx_N: (Mi,) # count self.ema_cluster_size.scatter_add_(0, idx, torch.ones_like(idx, dtype=self.ema_cluster_size.dtype)) # sum self.ema_embed_sum.scatter_add_(0, idx.unsqueeze(-1).expand(-1, C), rest_NC_b.detach() * self.ema_weights_list[si]) min_encoding_indices.append(idx) h_LC_b = h_NC_b[inv] h_BLC_list.append(h_LC_b) h_BLC = torch.stack(h_BLC_list, dim=0) # (B, L, C) h_BLC = self.phi[int(si/SN)](h_BLC) f_hat = f_hat + h_BLC f_rest -= h_BLC mean_vq_loss_i = F.mse_loss(f_hat.detach(), f_BLC).mul_(self.beta) + F.mse_loss(f_hat, f_no_grad) vq_loss_dict[f'vq_loss_{si}'] = mean_vq_loss_i.item() mean_vq_loss += mean_vq_loss_i min_encoding_stages_indices.append(min_encoding_indices) mean_vq_loss *= 1. / SN f_hat = f_hat.detach() - f_no_grad + f_BLC return f_hat, f_BLC, mean_vq_loss, min_encoding_stages_indices, vq_loss_dict, floatcoords_N2 def normalize_to_minus1_1(self, x): x_max_minus_min = x.max(dim=1, keepdim=True)[0] - x.min(dim=1, keepdim=True)[0] + 1e-8 max_val = x_max_minus_min.max(dim=-1, keepdim=True)[0] x_norm = x_max_minus_min / (max_val + 1e-8) x_normalized = ((x - x.min(dim=1, keepdim=True)[0]) / x_max_minus_min) * x_norm x_normalized = 2. * x_normalized - 1. # [-1, 1] return x_normalized def downsample(self, points): B, N, C = points.shape device = points.device out = [] for b in range(B): pc = points[b] # (N, 3) ratio = self.token_N2_nums / N idx = fps(pc, ratio=ratio) # (target_n,) pc_out = pc[idx] # (target_n, 3) out.append(pc_out) return torch.stack(out, dim=0) # (B, target_n, 3) def chamfer_distance(self, x, y): xx = torch.sum(x**2, dim=2) yy = torch.sum(y**2, dim=2) zz = torch.matmul(x, y.transpose(2, 1)) rx = xx.unsqueeze(2).expand(-1, -1, y.size(1)) ry = yy.unsqueeze(1).expand(-1, x.size(1), -1) P = rx + ry - 2*zz return torch.mean(torch.min(P, dim=2)[0]) + torch.mean(torch.min(P, dim=1)[0]) def get_vq_weight(self, step): vq_w_start = self.lambda_vq_start vq_w_end = self.lambda_vq_final if step >= self.warmup_steps: return vq_w_end # linear decay alpha = step / self.warmup_steps return vq_w_start *(1 - alpha) + vq_w_end * alpha def compute_usage_loss(self, eps=1e-12): """ Use EMA cluster_size to compute global usage entropy loss """ cluster_size = self.ema_cluster_size # (K,) p = cluster_size / (cluster_size.sum() + eps) entropy = - (p * (p + eps).log()).sum() # maximize entropy → minimize -entropy usage_loss = -entropy return usage_loss def forward(self, x, epoch, global_step): # voxelize point cloud voxel_feats, voxel_coords, spatial_shape, x_normalized = self.point_cloud_to_voxel(x) x_normalized = self.downsample(x_normalized) # encode & quantize f_hat, f_gt, vq_loss, token_label_codebook_idxs, vq_loss_dict, floatcoords = self.encode(voxel_feats, voxel_coords, spatial_shape, epoch) vq_weight = self.get_vq_weight(epoch) vq_loss = vq_weight * vq_loss # decode & reconstruct reconstructed = self.normalize_to_minus1_1(self.decoder(f_hat)) # Calculate Chamfer distance loss chamfer_loss = self.chamfer_distance(x_normalized, reconstructed) # [-1 1] chamfer_loss = self.lambda_chamfer * chamfer_loss # Calculate Construction loss construction_loss = F.smooth_l1_loss(reconstructed, floatcoords, reduction='mean') # [-1 1] construction_loss = self.lambda_construction * construction_loss # Calculate usage loss usage_loss = self.compute_usage_loss() usage_loss = self.usage_weight * usage_loss # Total loss is the sum of Chamfer distance and VQ losses total_loss = construction_loss + vq_loss + chamfer_loss + usage_loss # ema update codebook self.update_ema(self.ema_decay) # stateful controller (init once) current_vq_pct = vq_loss.detach() / (vq_loss.detach() + construction_loss.detach() + chamfer_loss.detach()) self.vq_stateful_controller(current_vq_pct, epoch) # print monitor information if global_step % 20 == 0: print("[",epoch,"/",global_step,"/",x.shape[0],"]") print("[Monitor] Token_feats(z_e) std: ", f_gt.std(dim=1).mean().item()) print("[Monitor] Recon_feats(z_q) std: ", f_hat.std(dim=1).mean().item()) active_codes = (self.ema_cluster_size > 1e-3).sum().item() print("[Monitor] Active embeddings:", active_codes, "/", self.codebook_size) print("[Monitor] Embedding weight std:",self.embeddings.weight.std().item()) # ===== Perplexity computation ===== probs = self.cluster_size / (self.cluster_size.sum() + 1e-10) # (K,) perplexity = torch.exp(-torch.sum(probs * torch.log(probs + 1e-10))) print("[Monitor] Perplexity:", perplexity.item()) return reconstructed, total_loss, construction_loss, vq_loss, chamfer_loss, usage_loss def vq_stateful_controller(self, current_vq_pct, epoch): if epoch < 100: return self.vq_window.append(current_vq_pct) # current_vq_pct = vq_loss / (recon + chamfer + vq) vq_pct_ma = sum(self.vq_window) / len(self.vq_window) if vq_pct_ma > self.safe_vq_pct_thr: self.patience_counter += 1 else: self.patience_counter = 0 if self.patience_counter >= 10 and self.triggers < 100: self.lambda_vq_final = max(self.lambda_vq_final * 0.99, self.lowest_vq_weight) self.triggers += 1 self.patience_counter = 0