（开源）深度学习推理引擎EasyTensor

2025-10-15

开源地址：https://github.com/JyAether/easytensor

项目概述

我认为不实现一遍都不算是真正理解本质， EasyTensor 是我最近一个基于 Python 的轻量级深度学习框架，旨在提供类似 PyTorch 的 API 设计，让深度学习变得更加简单易用。项目支持自动微分、多维数组运算、GPU 加速和智能内存管理，以及 Transformer，Bert 等主流模型也提供了完整的实现细节，同时由于深度学习计算图过于复杂，对于初学者，我提供了兼容层 V1 版本，基于原始Node实现，方便大家理解深度学习的计算逻辑。

核心特性

  • 自动微分: 完整的反向传播机制
  • GPU 加速: 基于 CuPy 的 CUDA 支持
  • 内存管理: 智能内存池和自动垃圾回收
  • 模块化设计: 高度可扩展的神经网络模块
  • 向后兼容: 保留 v1 版本接口，确保平滑升级

项目架构

目录结构总览

EasyTensor/
├── core/                    # 核心模块
│   ├── tensor.py           # 核心 Tensor 类
│   ├── device.py           # 设备管理
│   ├── model_io.py         # 模型保存和加载
│   ├── nn/                 # 神经网络模块
│   ├── data/               # 数据处理
│   ├── optim/              # 优化器
│   ├── utils/              # 工具模块
│   └── v1/                 # 兼容层
├── test/                   # 测试和示例
├── biz/                    # 业务示例
└── utils/                  # 绘图工具



EasyTensor/
├── core/                    # 核心模块
│   ├── tensor.py           # 核心Tensor类
│   ├── device.py           # 设备管理
│   ├── model_io.py         # 模型保存和加载
│   ├── nn/                 # 神经网络模块
│   │   ├── tensor_nn.py    # 基础神经网络层
│   │   ├── modules/        # 具体模块实现
│   │   │   ├── conv.py     # 卷积层
│   │   │   ├── embedding.py # 嵌入层
│   │   │   ├── pooling.py  # 池化层
│   │   │   └── rnn.py      # 循环神经网络
│   │   ├── attention.py    # 注意力机制
│   │   ├── bert_gpt.py     # BERT/GPT模型
│   │   ├── transform.py    # Transformer模型
│   │   └── distill.py      # 知识蒸馏
│   ├── data/               # 数据处理
│   │   ├── dataloader.py   # 数据加载器
│   │   └── word2vec.py     # 词向量
│   ├── optim/              # 优化器
│   │   └── lr_scheduler.py # 学习率调度器
│   ├── utils/              # 工具模块
│   │   ├── memory_utils.py # 内存管理
│   │   ├── tokenizer.py    # 分词器
│   │   └── serialization.py # 序列化工具
│   └── v1/                 # 兼容层
│       ├── engine.py       # 原始Node类
│       ├── nn.py           # 原始神经网络模块
│       └── optim/          # 原始优化器
├── test/                   # 测试和示例
│   ├── unit/               # 单元测试
│   ├── forward/            # 前向传播测试
│   └── network/            # 网络测试
└── biz/                    # 业务示例
    └── cnn.py              # CNN示例

核心模块详解

1. 核心 Tensor 类 (core/tensor.py)

功能概述: 深度学习的基本数据类型，支持多维张量操作和自动微分

主要特性:

  • 多维数组存储和管理
  • 自动微分和梯度计算
  • GPU/CPU 设备切换
  • 广播机制支持
  • 形状操作 (reshape, transpose, squeeze)
  • 数学运算 (加减乘除、矩阵乘法、激活函数)

关键方法:

  • backward(): 反向传播
  • cuda() / cpu(): 设备切换
  • reshape(), transpose(): 形状操作
  • sum(), mean(): 聚合操作

2. 神经网络模块 (core/nn/)

2.1 基础神经网络层 (tensor_nn.py)

模块基类: Module – 所有神经网络层的基类

支持的层类型:

  • Linear: 全连接层
  • ReLU, Sigmoid, Tanh: 激活函数层
  • BatchNorm1d: 批归一化
  • Dropout: 正则化层
  • Sequential: 顺序容器

损失函数:

  • MSELoss: 均方误差损失
  • CrossEntropyLoss: 交叉熵损失
  • BCEWithLogitsLoss: 二元交叉熵损失

优化器:

  • SGD: 随机梯度下降
  • Adam: 自适应矩估计

2.2 高级模块 (modules/)

卷积层 (conv.py):

  • Conv2d: 2D 卷积层，支持完整的梯度反向传播
  • 支持 stride, padding, bias 等参数

循环神经网络 (rnn.py):

  • RNN: 循环神经网络
  • LSTM: 长短期记忆网络

嵌入层 (embedding.py):

  • Embedding: 词嵌入层

池化层 (pooling.py):

  • MaxPool2d: 最大池化
  • AvgPool2d: 平均池化

2.3 高级功能

注意力机制 (attention.py):

  • Attention: 基础注意力机制
  • MultiHeadAttention: 多头注意力
  • SelfAttention: 自注意力

Transformer 模型 (transform.py):

  • PositionalEncoding: 位置编码
  • TransformerEncoderLayer: Transformer 编码器层

BERT/GPT 模型 (bert_gpt.py):

  • BERT: 双向编码器表示
  • GPT: 生成式预训练模型

知识蒸馏 (distill.py):

  • 模型压缩和知识迁移

3. 数据处理模块 (core/data/)

数据加载器 (dataloader.py):

  • 支持自定义数据集
  • 批量数据加载
  • 数据预处理管道

词向量 (word2vec.py):

  • Word2Vec 嵌入支持
  • 词汇表管理

4. 优化器模块 (core/optim/)

学习率调度器 (lr_scheduler.py):

  • StepLR: 阶梯学习率
  • ExponentialLR: 指数衰减
  • CosineAnnealingLR: 余弦退火

5. 工具模块 (core/utils/)

内存管理 (memory_utils.py):

  • MemoryPool: 内存池管理器
  • MemoryMonitor: 内存监控
  • memory_context: 内存上下文管理器
  • 支持 CPU 和 GPU 内存管理

分词器 (tokenizer.py):

  • 文本预处理和分词

序列化 (serialization.py):

  • 模型保存和加载

6. 兼容层 (core/v1/)

原始引擎 (engine.py):

  • Node: 原始的计算节点类
  • 保持向后兼容性

原始神经网络 (nn.py):

  • 旧版本的神经网络实现

原始优化器 (optim/):

  • SGD: 原始 SGD 实现
  • Adam: 原始 Adam 实现

测试模块结构

单元测试 (test/unit/)

基础功能测试:

  • demo_basic_tensor_operations.py: 基础张量操作演示
  • adam_test.py: Adam 优化器测试
  • momentum_test.py: 动量测试

网络测试:

  • deep_network_test.py: 深度网络测试
  • test_batch_norm.py: 批归一化测试
  • test_conv_layer.py: 卷积层测试
  • test_dropout_backward.py: Dropout 反向传播测试

损失函数测试:

  • test_loss.py: 损失函数测试

学习率调度测试:

  • test_step_lr.py: 阶梯学习率测试
  • train_with_step_lr_and_bn.py: 学习率调度和批归一化训练

多分类测试:

  • multi_class_network.py: 多分类网络测试

前向传播测试 (test/forward/)

  • forward_test.py: 前向传播功能测试
  • forward_test.ipynb: Jupyter 测试笔记本
  • torch_visible_test.ipynb: PyTorch 对比测试

网络测试 (test/network/)

  • custom_network.py: 自定义网络测试

其他测试

  • test_full_topo.py: 完整拓扑测试
  • test_topo_length.py: 拓扑长度测试
  • topo_analysis.py: 拓扑分析
  • 横向对比测试.py: 性能横向对比
  • 引擎测试.py: 引擎功能测试

业务示例

CNN 示例 (biz/cnn.py)

  • 卷积神经网络的完整实现示例
  • 图像分类任务演示

性能特点

支持的操作类型

  • 矩阵乘法、元素级运算
  • 激活函数 (ReLU, Sigmoid, Tanh, Softmax)
  • 卷积、池化、循环神经网络
  • 注意力机制、Transformer

内存管理特性

  • 智能内存池分配
  • 实时内存监控
  • 自动垃圾回收
  • GPU 内存优化

GPU 加速支持

  • 基于 CuPy 的 CUDA 支持
  • CPU/GPU 无缝切换
  • 混合精度计算

版本兼容性

v1 兼容层

项目保留了 v1 版本的接口，确保现有代码可以平滑升级：

  • core/v1/engine.py: 原始 Node 类
  • core/v1/nn.py: 原始神经网络模块
  • core/v1/optim/: 原始优化器实现

新版本优势

  • 更完善的 API 设计
  • 更好的性能优化
  • 更丰富的功能支持
  • 更强的内存管理

开发状态

已完成功能

✅ 核心 Tensor 类
✅ 基础神经网络层
✅ 卷积和循环神经网络
✅ 注意力机制和 Transformer
✅ BERT/GPT 模型
✅ 内存管理系统
✅ GPU 加速支持
✅ 完整的测试套件

持续优化

 性能优化
 内存效率提升
 API 完善
 文档补充

使用场景

教育用途

  • 深度学习原理学习
  • 自动微分机制理解
  • 神经网络实现细节

研究用途

  • 算法原型验证
  • 自定义层开发
  • 模型架构实验

轻量级应用

  • 简单深度学习任务
  • 资源受限环境
  • 快速原型开发

技术栈

核心依赖

  • NumPy: 数值计算基础
  • CuPy: GPU 加速支持
  • psutil: 系统资源监控

可选依赖

  • matplotlib: 可视化支持
  • scikit-learn: 数据预处理
  • jupyter: 交互式开发

总结

EasyTensor 是一个功能完整的轻量级深度学习框架，具有以下特点：

  1. 完整的深度学习功能: 从基础张量操作到高级 Transformer 模型
  2. 优秀的性能: 支持 GPU 加速和智能内存管理
  3. 良好的兼容性: 保留旧版本接口，支持平滑升级
  4. 丰富的测试: 全面的测试覆盖，确保功能正确性
  5. 清晰的架构: 模块化设计，易于理解和扩展

该项目既适合深度学习学习，也适合进行算法研究和原型开发。通过提供类似 PyTorch 的 API 设计，降低了学习成本，同时通过完整的实现细节，帮助用户深入理解深度学习框架的工作原理。在理解 GPT 等算法优化工作中可以通过该项目快速还原算法本质，比如FlashAttention检查点技术减少反向传播需要计算的特征图的数量是如何实现的等等

分类:作品
评论