Tez: 一个简单的 PyTorch 训练器

注意：目前我们不接受任何 Pull Request！所有 PR 都将被关闭。如果您需要某个功能或遇到问题，请提交 Issue。

tez（तेज़ / تیز）意为“锋利的、快速的、活跃的”。这是一个简单直接的库，旨在让您的 PyTorch 训练更加轻松。

该库目前仍处于早期阶段！因此，可能会有破坏性变更。

Tez 的核心理念很简单：

• 尽可能保持简单
• 提供尽可能高的可定制性
• 代码整洁
• 加速原型开发
• 适合生产环境

目前，tez 支持 CPU、单 GPU、多 GPU 以及 TPU 训练。更多功能即将推出！

使用 tez 非常简单。我们希望您不会远离 PyTorch 核心。因此，您可以完全自主地完成各项操作，只需借助 tez 来简化部分流程即可。

Tez 快速上手指南

Tez（意为“锐利、快速、活跃”）是一个极简的 PyTorch 训练库，旨在让模型训练更简单、代码更整洁，同时保持高度的可定制性。它支持 CPU、单 GPU、多 GPU 及 TPU 训练，适合快速原型开发和生产环境。

环境准备

• 系统要求：Linux / macOS / Windows
• Python 版本：建议 Python 3.7+
• 前置依赖：
  • PyTorch（需预先安装，Tez 不自动安装 PyTorch）
  • tqdm（用于进度条显示）

请确保已安装对应环境的 PyTorch。例如，使用 pip 安装 CUDA 版本：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

💡 国内用户可使用清华镜像加速 PyTorch 安装：

pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

通过 pip 直接安装 Tez：

pip install tez

💡 国内用户推荐使用清华镜像源加速安装：

pip install tez -i https://pypi.tuna.tsinghua.edu.cn/simple

⚠️ 注意：当前项目处于早期阶段，可能存在破坏性更新；暂不接受 Pull Request，如有问题请提交 Issue。

基本使用

Tez 的设计理念是“你掌控一切，Tez 帮你简化流程”。你只需定义模型、数据加载器和训练步骤，其余由 Tez 自动处理。

以下是最小可用示例：

import torch
import torch.nn as nn
from tez import Tez, TezArgs

# 1. 定义你的模型
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)
    
    def forward(self, x):
        return self.linear(x)

# 2. 准备数据集（此处以虚拟数据为例）
train_dataset = [(torch.randn(10), torch.randint(0, 2, (1,))) for _ in range(100)]
valid_dataset = [(torch.randn(10), torch.randint(0, 2, (1,))) for _ in range(20)]

# 3. 配置训练参数
args = TezArgs(epochs=5, batch_size=16, device="cuda" if torch.cuda.is_available() else "cpu")

# 4. 初始化并运行训练
model = MyModel()
trainer = Tez(model, args=args)
trainer.train(train_dataset, valid_dataset)

运行后，Tez 将自动处理训练循环、损失计算、进度显示及设备调度。你仍可完全控制模型结构、优化器、损失函数等细节，只需按需扩展即可。