orbit-torch 0.0.4a1__py3-none-any.whl

orbit/utils/mask.py

@@ -0,0 +1,92 @@
+import torch
+
+
+def make_padding_mask(src: torch.Tensor, pad_idx: int = 0) -> torch.Tensor:
+    '''
+    创建填充掩码。
+
+    Args:
+        src (torch.Tensor): 源序列张量。形状为 [B, L_src]
+        pad_idx (int, optional): 填充符号的索引。默认为 0。
+
+    Returns:
+        torch.Tensor: 填充掩码。形状为 [B, 1, 1, L_src]
+        True 表示该位置不是填充，应该被关注。
+    '''
+    mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
+    return mask
+
+
+def make_lookahead_mask(size: int, device: torch.device = torch.device('cpu')) -> torch.Tensor:
+    '''
+    创建前瞻掩码（下三角矩阵）。
+
+    Args:
+        size (int): 序列长度。
+        device (torch.device, optional): 设备。默认为 cpu。
+
+    Returns:
+        torch.Tensor: 前瞻掩码。形状为 [size, size]
+        True 表示允许关注的位置（下三角部分）。
+    '''
+    mask = torch.tril(torch.ones((size, size), device=device)).bool()
+    return mask
+
+
+def make_causal_mask(tgt: torch.Tensor, pad_idx: int = 0) -> torch.Tensor:
+    '''
+    创建因果掩码（结合了填充掩码和前瞻掩码）。
+
+    Args:
+        tgt (torch.Tensor): 目标序列张量。形状为 [B, L_tgt]
+        pad_idx (int, optional): 填充符号的索引。默认为 0。
+
+    Returns:
+        torch.Tensor: 因果掩码。形状为 [B, 1, L_tgt, L_tgt]
+    '''
+    pad_mask = make_padding_mask(tgt, pad_idx)
+    seq_len = tgt.size(1)
+    lookahead_mask = make_lookahead_mask(seq_len, device=tgt.device)
+
+    # pad_mask: [B, 1, 1, L]
+    # lookahead_mask: [L, L]
+    # 广播后: [B, 1, L, L]
+    mask = pad_mask & lookahead_mask
+    return mask
+
+
+def make_sliding_window_mask(
+    tensor: torch.Tensor, window_size: int, pad_idx: int = 0, causal: bool = True
+) -> torch.Tensor:
+    '''
+    创建滑动窗口掩码。
+
+    Args:
+        tensor (torch.Tensor): 输入序列张量。形状为 [B, L]
+        window_size (int): 窗口大小（单侧）。
+        pad_idx (int, optional): 填充符号的索引。默认为 0。
+        causal (bool, optional): 是否为因果（单向）。默认为 True。
+            如果为 True，位置 i 只能关注 [i - window_size, i]。
+            如果为 False，位置 i 可以关注 [i - window_size, i + window_size]。
+
+    Returns:
+        torch.Tensor: 滑动窗口掩码。形状为 [B, 1, L, L]
+    '''
+    pad_mask = make_padding_mask(tensor, pad_idx)  # [B, 1, 1, L]
+    seq_len = tensor.size(1)
+
+    ones = torch.ones((seq_len, seq_len), device=tensor.device, dtype=torch.bool)
+
+    if causal:
+        # j <= i AND j >= i - window_size
+        window_mask = torch.tril(ones, diagonal=0) & torch.triu(
+            ones, diagonal=-window_size
+        )
+    else:
+        # j <= i + window_size AND j >= i - window_size
+        window_mask = torch.tril(ones, diagonal=window_size) & torch.triu(
+            ones, diagonal=-window_size
+        )
+
+    mask = pad_mask & window_mask
+    return mask

orbit/utils/seed.py

@@ -0,0 +1,66 @@
+import torch
+import torch
+import numpy as np
+import random
+import os
+
+def seed_everything(seed=42, strict=False):
+    """
+    设置所有随机种子以确保 PyTorch 实验的可复现性。
+    
+    Args:
+        seed (int): 随机种子数值，默认为 42。
+        strict (bool): 是否启用严格确定性模式。
+                       如果为 True，将调用 torch.use_deterministic_algorithms(True)，
+                       这可能会导致某些不支持确定性算法的操作报错，并且会降低训练速度。
+    """
+    import orbit
+    orbit.seed_info = seed
+    # 1. 设置 Python 原生 random
+    random.seed(seed)
+    
+    # 2. 设置 Python 哈希种子 (影响字典/集合迭代顺序)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    
+    # 3. 设置 Numpy
+    np.random.seed(seed)
+    
+    # 4. 设置 PyTorch CPU/GPU
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    
+    # 5. 设置 CuDNN 后端 (常规复现性设置)
+    if torch.cuda.is_available():
+        # 禁止寻找最优算法 (因为最优算法可能因硬件状态而变)
+        torch.backends.cudnn.benchmark = False 
+        # 强制使用确定性算法
+        torch.backends.cudnn.deterministic = True 
+    # 6. 严格模式 (Strict Mode)
+    if strict:
+        try:
+            # 启用严格确定性算法
+            # 注意：某些操作如果 PyTorch 没有对应的确定性实现，会直接通过 RuntimeError 报错
+            torch.use_deterministic_algorithms(True)
+            
+            # 为了让 use_deterministic_algorithms 在 CUDA 上正常工作，
+            # 必须设置 CUBLAS_WORKSPACE_CONFIG，否则会报 CuBLAS 错误。
+            # :4096:8 是官方推荐的设置，虽然会增加少许显存开销
+            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+            
+            print(f"[Info] Strict deterministic mode enabled. (seed={seed})")
+        except AttributeError:
+            print("[Warning] torch.use_deterministic_algorithms is not available in your PyTorch version.")
+    else:
+        print(f"[Info] Random seed set as {seed}")
+
+def worker_init_fn(worker_id):
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+def create_generator() -> torch.Generator:
+    """创建随机数生成器"""
+    import orbit
+    seed = orbit.seed_info if hasattr(orbit, 'seed_info') else 42
+    return torch.Generator().manual_seed(seed)

orbit_torch-0.0.4a1.dist-info/METADATA

@@ -0,0 +1,25 @@
+Metadata-Version: 2.4
+Name: orbit-torch
+Version: 0.0.4a1
+Summary: A PyTorch training engine with plugin system
+Author: Aiden Hopkins
+Author-email: acdphc@qq.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Requires-Dist: torch>=1.10.0
+Requires-Dist: rich
+Requires-Dist: tensorboard
+Requires-Dist: matplotlib
+Requires-Dist: seaborn
+Requires-Dist: numpy
+Requires-Dist: scikit-learn
+Requires-Dist: einops
+Requires-Dist: tokenizers
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

orbit_torch-0.0.4a1.dist-info/RECORD

@@ -0,0 +1,29 @@
+orbit/__init__.py,sha256=xikJCuBzK0VSwJ8gvPt0cZijw5TzVI2t3a62_aK-uEo,51
+orbit/callback.py,sha256=FXq-bOVfoYW0S5S95ry55yCn1QgYkecpHonO35BNKwE,1738
+orbit/engine.py,sha256=4Vskd39cfVa7jUWucikva8oDc68om4DtGH6OZLUDhKE,34230
+orbit/optim/__init__.py,sha256=jp3TZFLM5LRcHeeYjR6qNwzwo1dfeR65A6DSGCEPnFg,67
+orbit/optim/muon.py,sha256=IT0l1b0mAtLpB-SUuPyo9uaYHO2x_wkqg1TLcDJ6Bvk,7342
+orbit/optim/sam.py,sha256=lDfou1jQYM_qtlWgYzfUKp2Zdi-uBSKZS7iYxB-yGkU,3702
+orbit/plugin/__init__.py,sha256=9EcMoS__kzldM4MBNUwv5xUs21rliqrBALVmzIk_ADw,435
+orbit/plugin/board.py,sha256=dK3aaTmznq9vSYvpSPgfE9kt30tieR93oK0oDaC8yJw,2267
+orbit/plugin/checkpoint.py,sha256=w9R5O2-jzbRPSGPiuGuUycq21EebAdqigEFK0T0waOE,10898
+orbit/plugin/classification.py,sha256=iVOFNUaL4SUMNEtnmbnioprb2nH1sOV4IsGH7UPhU8U,7019
+orbit/plugin/display_model.py,sha256=Uch563Jq0R78lGHYiDPotXB6dmWfqdClSOiuOIKw7i8,2907
+orbit/plugin/early_stopping.py,sha256=eDTjIzxDZvnCXANR2RXn8vrMh5wX8F2rQjl0BNp-6dU,4065
+orbit/plugin/ema.py,sha256=NtQVB3Rz5YreQwWicXCNy_CfD6d7pu0crGGrz6-IJI8,3892
+orbit/plugin/gradient_accumulation.py,sha256=JaSl-U9Zo_jeEaJE-KAXEPhQdjFA_He729VaCw2W8n4,1066
+orbit/plugin/memory_estimator.py,sha256=YxSXbjAwl6jJh6yFEod5Tu_pgZh1Qadb0JsgmyfF7Mc,9743
+orbit/plugin/mentor.py,sha256=sx3gLHnxYJX1tJhqnqHPTW2KuTeNM7h0PqOIquxMaxQ,12474
+orbit/plugin/overfit.py,sha256=lpfxX4igi4uNdY4OR2_fpayw52iEfkCQnQ0HgzodZd4,1194
+orbit/plugin/warmup.py,sha256=jS3Q52zuhTdqN9Bw85zcoH1YXZoqyRvnD7JHSSLRUwQ,4996
+orbit/plugin/data/mentor_i18n.json,sha256=KHnEKRKHeLga_nqEUZVgiTh4jb4j2f_yD0MJT-7sdis,8005
+orbit/utils/__init__.py,sha256=gXNfTz1a_kL5UZXT2o5BhoItc4Dgrk8w4K2hsfkbiYk,589
+orbit/utils/freeze.py,sha256=ujR5wl2GHuqIQi0rm2i_yz6W8dj6X4YqgyZ3tBbv4yU,2405
+orbit/utils/initialization.py,sha256=X0-OP0FPWX60_SRuB-x_Od3KdDh4EAD9UbOCZFj6NtI,20560
+orbit/utils/layer_io.py,sha256=E-YnYw7av2ZdczDzhs08hUQtbZrcHPvqDdzEvWgs7Gs,1872
+orbit/utils/mask.py,sha256=DayMxmKlWMYMYPceYhYlf9lIx4FD9s4Tp1WvzfVpsu0,3072
+orbit/utils/seed.py,sha256=tbjF2jIPQfb0M_-wXOMVH3_XzWVsDkuhzD5lctGqalA,2578
+orbit_torch-0.0.4a1.dist-info/METADATA,sha256=g2NSH_PVm9afSyVAfvct5xvk2hFColq85tKdbGQfNFc,700
+orbit_torch-0.0.4a1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+orbit_torch-0.0.4a1.dist-info/top_level.txt,sha256=emrF0of931NzTSL4R5yBKpGoewFCB-cAwYNcUF5cqBs,6
+orbit_torch-0.0.4a1.dist-info/RECORD,,

orbit_torch-0.0.4a1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

orbit_torch-0.0.4a1.dist-info/top_level.txt

@@ -0,0 +1 @@
+orbit