Add code/cube3d/training/utils.py

code/cube3d/training/utils.py

@@ -0,0 +1,341 @@
+import os
+import logging
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig, OmegaConf
+from safetensors.torch import load_model, save_model
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+from matplotlib.ticker import MaxNLocator
+
+BOUNDING_BOX_MAX_SIZE = 1.925
+
+
+def normalize_bbox(bounding_box_xyz: Tuple[float]):
+    #import ipdb; ipdb.set_trace()
+    max_l = max(bounding_box_xyz)
+    return [BOUNDING_BOX_MAX_SIZE * elem / max_l for elem in bounding_box_xyz]
+
+def normalize_bboxs(bounding_box_xyz, max_xyz):
+    #max_l = max(bounding_box_xyz)
+    normalized = BOUNDING_BOX_MAX_SIZE * bounding_box_xyz / torch.tensor(max_xyz, device=bounding_box_xyz.device)
+    return normalized
+
+def load_config(cfg_path: str) -> Any:
+    """
+    Load and resolve a configuration file.
+    Args:
+        cfg_path (str): The path to the configuration file.
+    Returns:
+        Any: The loaded and resolved configuration object.
+    Raises:
+        AssertionError: If the loaded configuration is not an instance of DictConfig.
+    """
+
+    cfg = OmegaConf.load(cfg_path)
+    OmegaConf.resolve(cfg)
+    assert isinstance(cfg, DictConfig)
+    return cfg
+
+
+def parse_structured(cfg_type: Any, cfg: DictConfig) -> Any:
+    """
+    Parses a configuration dictionary into a structured configuration object.
+    Args:
+        cfg_type (Any): The type of the structured configuration object.
+        cfg (DictConfig): The configuration dictionary to be parsed.
+    Returns:
+        Any: The structured configuration object created from the dictionary.
+    """
+
+    scfg = OmegaConf.structured(cfg_type(**cfg))
+    return scfg
+
+
+def load_model_weights(model: torch.nn.Module, ckpt_path: str) -> None:
+    """
+    Load a safetensors checkpoint into a PyTorch model.
+    The model is updated in place.
+
+    Args:
+        model: PyTorch model to load weights into
+        ckpt_path: Path to the safetensors checkpoint file
+
+    Returns:
+        None
+    """
+    assert ckpt_path.endswith(
+        ".safetensors"
+    ), f"Checkpoint path '{ckpt_path}' is not a safetensors file"
+
+    load_model(model, ckpt_path)
+
+
+def save_model_weights(model: torch.nn.Module, save_path: str) -> None:
+    """
+    Save a PyTorch model as safetensors format.
+    
+    Args:
+        model: PyTorch model to save
+        save_path: Output path (must end with .safetensors)
+    """
+    assert save_path.endswith(".safetensors"), "Path must end with .safetensors"
+    
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    
+    save_model(model, save_path)
+    
+    assert os.path.exists(save_path), f"Failed to save to {save_path}"
+
+
+def select_device() -> Any:
+    """
+    Selects the appropriate PyTorch device for tensor allocation.
+
+    Returns:
+        Any: The `torch.device` object.
+    """
+    return torch.device(
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps"
+        if torch.backends.mps.is_available()
+        else "cpu"
+    )
+
+def mask_cross_entropy(p_st, p_ed, p_max, logits, target, shift):
+    p_range = torch.arange(p_st, p_ed, device=logits.device)
+    p_range_expanded = p_range.unsqueeze(0).repeat(p_max.shape[0], 1)
+    valid_p_mask = p_range_expanded <= p_max.unsqueeze(1)+p_st
+
+    valid_p_mask = valid_p_mask.unsqueeze(1).expand(-1, logits.shape[1], -1)
+    logits_masked = logits.clone()
+    logits_masked[:,:,p_st:p_ed][~valid_p_mask] = float('-inf')
+
+    p_loss = F.cross_entropy(
+        logits_masked[:, :-1, p_st:p_ed].permute(0, 2, 1), 
+        target[:, shift:, p_st:p_ed].argmax(-1), 
+    )
+
+    return p_loss
+
+def positional_encoding(x, num_freqs):
+    
+    freqs = 2.0 ** torch.arange(num_freqs, device=x.device)  # [num_freqs]
+    angles = x.unsqueeze(-1) * freqs  # [..., num_freqs]
+    sin_cos = torch.cat([angles.sin(), angles.cos()], dim=-1)  # [..., 2*num_freqs]
+    return sin_cos.flatten(-2)  
+
+def visualize_token_probabilities(
+    probs, 
+    cut_idx, 
+    sample_idx=0, 
+    tokens_per_page=10,  # 每页显示的token数量
+    figsize=(12, 20),   # 单页图表大小
+    save_dir=None       # 保存图片的目录（None则直接显示）
+):
+    """
+    分页展示所有有效token的概率分布（每页10个，一行一个token）
+    
+    参数:
+    - probs: 概率张量，形状为 (batch_size, seq_len, num_classes)
+    - cut_idx: 有效区域的截止索引
+    - sample_idx: 要可视化的batch样本索引
+    - tokens_per_page: 每页显示的token数量
+    - figsize: 单页图表大小
+    - save_dir: 保存图片的目录（若为None则直接显示）
+    """
+    # 转换为numpy数组
+    if isinstance(probs, torch.Tensor):
+        probs = probs.cpu().detach().numpy()
+    
+    # 获取单个样本的概率分布
+    sample_probs = probs[sample_idx]  # (seq_len, num_classes)
+    seq_len, num_classes = sample_probs.shape
+    
+    # 处理cut_idx，确定有效区域并提取有效token
+    if isinstance(cut_idx, torch.Tensor):
+        cut_idx = cut_idx.cpu().detach().numpy()
+    valid_length = min(int(cut_idx[sample_idx] if not np.isscalar(cut_idx) else cut_idx), seq_len)
+    valid_probs = sample_probs[:valid_length, :]  # 只取有效区域内的token
+    num_valid_tokens = valid_probs.shape[0]
+    
+    if num_valid_tokens == 0:
+        print(f"警告：没有有效token可显示（有效区域长度：{valid_length}）")
+        return None
+    
+    # 创建保存目录（如果需要）
+    if save_dir is not None and not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    
+    # 计算总页数
+    total_pages = (num_valid_tokens + tokens_per_page - 1) // tokens_per_page
+    print(f"共{num_valid_tokens}个有效token，分为{total_pages}页展示")
+    
+    # 分页生成图表
+    figures = []
+    for page in range(total_pages):
+        # 计算当前页的token范围
+        start = page * tokens_per_page
+        end = min(start + tokens_per_page, num_valid_tokens)
+        page_tokens = end - start
+        
+        # 创建当前页的画布
+        fig, axes = plt.subplots(page_tokens, 1, figsize=(figsize[0], 2*page_tokens))
+        fig.suptitle(
+            f'Token Probability Distributions (Sample {sample_idx}) - Page {page+1}/{total_pages}',
+            fontsize=16,
+            y=1.02
+        )
+        
+        # 为当前页的每个token绘制分布
+        for i in range(page_tokens):
+            token_idx = start + i
+            token_probs = valid_probs[i]  # 当前页内的相对索引
+            ax = axes[i] if page_tokens > 1 else axes  # 处理单token情况
+            
+            # 绘制条形图
+            class_indices = np.arange(num_classes)
+            bars = ax.bar(class_indices, token_probs, width=0.8, color='skyblue', edgecolor='black')
+            
+            # 突出显示最高概率的类别
+            max_prob_idx = np.argmax(token_probs)
+            max_prob_value = token_probs[max_prob_idx]
+            bars[max_prob_idx].set_color('orange')
+            
+            # 标注概率>5%的类别
+            for j, (bar, prob) in enumerate(zip(bars, token_probs)):
+                height = bar.get_height()
+                if prob > 0.05:
+                    ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,
+                            f'{prob:.2f}', ha='center', va='bottom', fontsize=9)
+            
+            # 设置子图标题和坐标轴
+            ax.set_title(
+                f'Token {token_idx} (Max: Class {max_prob_idx} = {max_prob_value:.2f})',
+                fontsize=11
+            )
+            ax.set_xlabel('Class Index')
+            ax.set_ylabel('Probability')
+            ax.set_ylim(0, 1.1)
+            ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+            ax.grid(True, alpha=0.3, axis='y')
+            
+            # 除最后一个子图外隐藏x轴标签
+            if i != page_tokens - 1:
+                ax.set_xlabel('')
+        
+        plt.tight_layout()
+        figures.append(fig)
+        
+        # 保存或显示图表
+        if save_dir is not None:
+            save_path = os.path.join(save_dir, f'token_probs_page_{page+1}.png')
+            fig.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"已保存第{page+1}页至: {save_path}")
+        else:
+            plt.show()
+        plt.close(fig)  # 关闭当前页图表，释放内存
+    
+    return figures
+
+def visualize_max_prob_distribution(
+    probs, 
+    cut_idx=None,  # 不再需要，因为已提前过滤
+    sample_idx=0, 
+    bins=20,
+    figsize=(12, 6)
+):
+    # 转换为numpy数组
+    if isinstance(probs, torch.Tensor):
+        probs = probs.cpu().detach().numpy()
+    
+    # 获取单个样本的概率分布并计算最大概率
+    sample_probs = probs[sample_idx]
+    max_probs_per_token = np.max(sample_probs, axis=1)  # 所有token都是已过滤的有效token
+    
+    # 创建画布
+    fig, ax = plt.subplots(figsize=figsize)
+    
+    # 绘制直方图
+    n, bins, patches = ax.hist(
+        max_probs_per_token,
+        bins=bins,
+        range=(0, 1),
+        edgecolor='black',
+        alpha=0.7,
+        color='skyblue'
+    )
+    
+    # 标注数量
+    for count, patch in zip(n, patches):
+        height = patch.get_height()
+        if height > 0:
+            ax.text(
+                patch.get_x() + patch.get_width()/2., 
+                height + 0.5,
+                f'{int(count)}', 
+                ha='center', 
+                va='bottom', 
+                fontsize=9
+            )
+    
+    # 统计指标
+    mean_prob = np.mean(max_probs_per_token)
+    median_prob = np.median(max_probs_per_token)
+    max_count = int(np.max(n)) if len(n) > 0 else 0
+    
+    # 设置标题和坐标轴
+    ax.set_title(
+        f'Distribution of Maximum Probabilities (All Valid Tokens from 5 Iterations)\n'
+        f'Total tokens: {len(max_probs_per_token)} | Mean: {mean_prob:.2f} | Median: {median_prob:.2f}',
+        fontsize=14
+    )
+    ax.set_xlabel('Maximum Probability Value (0-1)')
+    ax.set_ylabel('Number of Tokens (Frequency)')
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, max_count + 2)
+    ax.xaxis.set_major_locator(MaxNLocator(nbins=11))
+    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
+    ax.grid(True, alpha=0.3, axis='y')
+    
+    plt.tight_layout()
+    return fig
+
+
+def top_k_prob_mask(probs, cut_idx, top_percent=0.15, visualize=False):
+    max_probs = probs.permute(0, 2, 1).max(dim=1).values  # (batch_size, seq_len)
+    batch_size, seq_len = max_probs.shape
+
+    # 1. 生成基础mask：cut_idx前面为True，后面为False
+    if isinstance(cut_idx, (int, float)):
+        cut_idx = torch.tensor([cut_idx] * batch_size, device=max_probs.device)
+    base_mask = (torch.arange(seq_len, device=max_probs.device)[None, :] < cut_idx[:, None])
+    valid_count = base_mask.sum().item()
+
+    # 处理无有效位置的情况
+    if valid_count == 0:
+        empty_mask = torch.zeros_like(max_probs, dtype=torch.bool)
+        return empty_mask, empty_mask
+
+    # 2. 计算原始目标mask（cut内前N%高概率True）
+    valid_probs = max_probs[base_mask]
+    total_valid = valid_probs.numel()
+    k = max(min(int(total_valid * top_percent), total_valid), 1)
+    _, top_valid_indices = torch.topk(valid_probs, k)
+
+    # 原始mask：cut内top k为True，其余全False
+    valid_area_original = torch.zeros(total_valid, dtype=torch.bool, device=max_probs.device)
+    valid_area_original[top_valid_indices] = True
+    original_mask = torch.zeros_like(max_probs, dtype=torch.bool)
+    original_mask[base_mask] = valid_area_original
+
+    # 3. 计算反向mask（cut内非top k为True，cut外全False）
+    valid_area_reverse = ~valid_area_original  # 与原始有效区域完全相反
+    reverse_mask = torch.zeros_like(max_probs, dtype=torch.bool)
+    reverse_mask[base_mask] = valid_area_reverse  # cut外保持False
+
+    return original_mask, reverse_mask  # 返回两个mask