TinyDL - 轻量级深度学习框架

📖 项目简介

最新版本： https://github.com/Leavesfly/TinyAI

TinyDL 是一个用 Java 实现的轻量级深度学习框架，旨在为深度学习初学者和研究人员提供清晰、简洁的核心功能实现。该框架参考了 PyTorch 的设计理念，实现了自动微分、神经网络层、优化器等核心组件，特别适合：

• 🎓 深度学习教学与学习：代码结构清晰，中文注释详尽，便于理解底层原理
• 🔬 学术研究与实验：模块化设计，易于扩展和定制
• 🚀 快速原型开发：提供完整的机器学习工具链，支持多种AI应用场景
• 💡 算法验证：在JVM环境中进行深度学习算法的验证和调试
• 🤖 多模态AI应用：支持NLP、CV、强化学习等多个AI领域

✨ 核心特性

🔢 多维数组计算

• NdArray 核心类：支持标量、向量、矩阵及高维张量操作，完整序列化支持
• 丰富数学运算：四则运算、矩阵乘法、形状变换、广播机制
• 内存高效：针对 CPU 优化的数组实现，支持缓存机制

🔄 自动微分系统

• 动态计算图：运行时构建，支持复杂的控制流
• 双模式反向传播：递归和迭代两种实现，避免栈溢出
• 自动梯度计算：一键调用 backward() 完成反向传播
• 灵活的梯度控制：支持梯度开关和计算图切断

🧠 神经网络组件

• 完整的层实现：
• 全连接层（LinearLayer、AffineLayer）
• 优化卷积层（ConvLayer支持偏置、Xavier初始化）
• 高级卷积层（BatchNormLayer、DepthwiseSeparableConvLayer）
• 循环层（LstmLayer、SimpleRnnLayer）
• Transformer组件（MultiHeadAttention、GPT2Block）
• 嵌入层（Embedding、GPT2TokenEmbedding）
• 激活函数（ReLU、Sigmoid、Tanh、Softmax）
• 模块化设计：Layer 和 Block 支持灵活组合和残差连接
• 预构建网络：MlpBlock、LstmBlock、SequentialBlock、GPT2Model

📊 完整机器学习工具链

• 数据处理：多种内置数据集（MNIST、螺旋数据、Word2Vec专用数据集等）
• 损失函数：交叉熵、均方误差、掩码损失等
• 优化算法：SGD、Adam 优化器，支持学习率调度
• 训练控制：Trainer 类提供完整训练循环和并行训练支持
• 效果评估：准确率、回归误差等评估器
• 模型管理：完整的序列化系统，支持检查点和压缩存储

🚀 并行训练系统

• 多线程训练：自动检测模型并行性，智能线程数分配
• 梯度聚合：支持并行批次处理和梯度累积，完整的GradientAggregator实现
• 资源管理：完整的线程池管理和异常处理，包含ParallelBatchProcessor
• 性能监控：并行训练统计和性能分析，支持训练性能对比

🤖 自然语言处理

• Word2Vec实现：支持Skip-gram和CBOW两种模式，完整词汇管理
• 负采样优化：基于词频的高效负采样算法，专用Word2VecDataSet
• GPT-2模型：完整的小规模语言模型实现，支持Token嵌入和位置编码
• MoE架构：混合专家模型，包含门控网络和专家网络，支持动态专家选择
• Transformer组件：多头注意力、位置编码、层归一化等完整实现
• 文本生成：支持自回归文本生成和下一个token预测
• 词向量操作：相似度计算、最相似词查找、词向量可视化等

🖼️ 计算机视觉增强

• CNN深度优化：Im2Col/Col2Im缓存机制，性能大幅提升
• 高级卷积技术：深度可分离卷积、批量归一化
• 灵活网络配置：SimpleConvNet支持残差连接和自定义架构
• 性能基准测试：完整的CNN性能分析和对比工具

📈 可视化与监控

• 训练监控：实时显示损失和准确率变化，支持文件日志
• 结果可视化：基于 JFreeChart 的图表绘制
• 模型结构图：UML 工具可视化网络架构
• 性能分析：并行训练效率统计和资源使用监控
• 强化学习可视化：支持RL训练过程可视化和性能分析

🎯 丰富应用示例

• 分类任务：螺旋数据分类、手写数字识别
• 回归任务：曲线拟合、时间序列预测
• 序列建模：RNN 序列预测、LSTM 应用、Seq2Seq架构
• 自然语言处理：Word2Vec训练、GPT-2文本生成、MoE-GPT实现
• 计算机视觉：卷积网络优化、深度可分离卷积、批量归一化
• Transformer应用：多头注意力、位置编码、Transformer编码器
• 强化学习：DQN智能体、REINFORCE算法、多臂老虎机问题、CartPole和GridWorld环境
• 并行训练：多线程训练演示和性能对比
• 模型序列化：完整的模型保存和加载示例

🎮 强化学习模块

• 核心算法：深度Q网络(DQN)、策略梯度(REINFORCE)、多臂老虎机算法
• 环境支持：CartPole倒立摆、GridWorld网格世界、自定义环境接口
• 训练工具：经验回放、目标网络、ε-贪婪策略、负载均衡
• 智能体管理：完整的Agent抽象和具体实现
• 性能评估：实时监控训练指标和性能分析

🏗️ 系统架构

TinyDL 采用分层模块化架构，各组件职责明确，易于理解和扩展：

graph TB
    subgraph "应用层"
        E[example包

示例程序] --> M[modality包

应用领域]
    end

    subgraph "机器学习层"
        M --> ML[mlearning包

训练/推理组件]
    end

    subgraph "神经网络层"
        ML --> N[nnet包

网络层和块]
    end

    subgraph "计算层"
        N --> F[func包

自动微分]
        F --> ND[ndarr包

多维数组]
    end

    style E fill:#fff2cc
    style M fill:#d5e8d4
    style ML fill:#dae8fc
    style N fill:#f8cecc
    style F fill:#e1d5e7
    style ND fill:#ffcce6

📦 核心模块详解

1. ndarr包 - 数值计算基础

• NdArray：多维数组核心实现，支持各种数学运算和序列化
• Shape：形状定义和操作，支持动态维度
• NdArrayUtil：数组工具方法和优化函数
• 设计理念：提供高效的CPU计算支持，为上层提供数值计算基础

2. func包 - 自动微分引擎

• Variable：变量抽象，记录计算图节点，支持迭代和递归反向传播
• Function：所有数学运算的基类
• 运算类别：
• base/：四则运算（Add、Sub、Mul、Div）
• math/：数学函数（Sin、Exp、ReLU、Sigmoid等）
• matrix/：矩阵运算（MatMul、Reshape、Softmax等）
• loss/：损失函数（MeanSE、SoftmaxCE等）
• 设计理念：通过计算图实现动态自动微分，避免栈溢出

3. nnet包 - 神经网络构建

• Layer：网络层接口，支持参数管理
• Block：网络块抽象，可组合多个层
• 层实现：
• layer/dnn/：全连接层（LinearLayer、AffineLayer）
• layer/cnn/：卷积层、池化层、BatchNorm、DepthwiseSeparableConv
• layer/rnn/：循环神经网络层（SimpleRnn、LSTM）
• layer/transformer/：Transformer组件（多头注意力、位置编码、GPT-2）
• layer/activate/：激活函数层
• layer/embedding/：嵌入层
• 块实现：
• block/：SequentialBlock、MlpBlock、LstmBlock
• block/transformer/：GPT2Block、TransformerEncoder等
• 设计理念：模块化组件，支持复杂的网络架构构建

4. mlearning包 - 机器学习工具链

• Model：模型封装器，支持序列化和模型信息管理
• Trainer：训练控制器，支持并行训练和简化版实现
• DataSet：数据集抽象和实现，包括Word2VecDataSet等专用数据集
• Loss：损失函数集合（交叉熵、均方误差等）
• Optimizer：优化器实现（SGD、Adam）
• Evaluator：模型评估器和准确率计算
• ModelSerializer：完整的模型序列化系统
• ParameterManager：参数管理和操作工具
• Monitor：训练过程监控和可视化
• 并行训练：parallel/包提供多线程训练支持
• 设计理念：提供企业级的机器学习开发工具链

5. modality包 - 领域应用

• cv/：计算机视觉应用
• SimpleConvNet：增强的卷积神经网络，支持残差连接
• nlp/：自然语言处理应用
• Word2Vec：完整的词向量实现（Skip-gram/CBOW）
• GPT2Model：小规模GPT-2语言模型
• 设计理念：针对特定领域的高层封装和优化

6. example包 - 实用示例

• classify/：分类任务示例（螺旋数据、MNIST）
• regress/：回归任务示例（曲线拟合、RNN预测）
• nlp/：自然语言处理示例（Word2Vec、GPT-2）
• cv/：计算机视觉示例（卷积网络优化）
• transformer/：Transformer相关示例
• parallel/：并行训练示例
• embedd/：嵌入层示例
• 设计理念：展示框架最新功能，提供学习参考

🚀 快速开始

环境要求

• Java 8+
• Maven 3.6+

依赖配置

<dependencies>
    <dependency>
        <groupId>jfree</groupId>
        <artifactId>jfreechart</artifactId>
        <version>1.0.7</version>
    </dependency>
    <dependency>
        <groupId>junit</groupId>
        <artifactId>junit</artifactId>
        <version>4.13.2</version>
        <scope>test</scope>
    </dependency>
</dependencies>

💡 快速上手示例

1. 自动微分基础

理解TinyDL的核心——自动微分机制：

// 创建变量（支持梯度计算）
Variable x = new Variable(new NdArray(2.0f)).setName("x");
Variable y = new Variable(new NdArray(3.0f)).setName("y");

// 构建计算表达式 z = (x + y) * x = (2 + 3) * 2 = 10
Variable z = x.add(y).mul(x);

// 自动微分：计算 dz/dx 和 dz/dy
z.backward();

// 查看梯度结果
System.out.println("z的值: " + z.getValue().getNumber()); // 输出: 10.0
System.out.println("x的梯度 dz/dx: " + x.getGrad().getNumber()); // 输出: 5.0 (y + x)
System.out.println("y的梯度 dz/dy: " + y.getGrad().getNumber()); // 输出: 2.0 (x)

2. 神经网络构建

使用Block构建多层感知机：

// 网络参数设置
int batchSize = 32;
int inputSize = 2;   // 输入维度
int hiddenSize = 10; // 隐藏层大小
int outputSize = 3;  // 输出类别数

// 创建多层感知机：input -> hidden -> output
MlpBlock mlpBlock = new MlpBlock("MLP", batchSize, null,
        inputSize, hiddenSize, outputSize);
Model model = new Model("ClassificationModel", mlpBlock);

// 创建随机输入数据
Variable input = new Variable(
        NdArray.likeRandom(-1, 1, new Shape(batchSize, inputSize))
);

// 前向传播
Variable output = model.forward(input);
System.out.println("输出形状: " + output.getValue().getShape()); // [32, 3]

3. 并行训练演示

使用多线程加速训练过程：

// 训练参数
int maxEpoch = 100;
int batchSize = 10;
float learningRate = 0.01f;
int threadCount = 4; // 并行线程数

// 创建数据集
SpiralDateSet dataSet = new SpiralDateSet(batchSize);

// 创建模型
MlpBlock block = new MlpBlock("ParallelMLP", batchSize, null, 2, 30, 3);
Model model = new Model("ParallelClassifier", block);

// 配置组件
Optimizer optimizer = new Adam(model, learningRate);
Loss lossFunc = new SoftmaxCrossEntropy();
Monitor monitor = new Monitor();
AccuracyEval evaluator = new AccuracyEval(new Classify(), model, dataSet);

// 创建并行训练器
Trainer trainer = new Trainer(maxEpoch, monitor, evaluator, true, threadCount);
trainer.init(dataSet, model, lossFunc, optimizer);

// 开始并行训练
trainer.parallelTrain(true); // true表示打乱数据

4. MoE混合专家模型

实现完整的MoE (Mixture of Experts) 架构：

// 创建MoE-GPT模型
MoEGPTModel model = MoEGPTModel.createSmallModel("demo_moe_gpt", vocabSize);

// 模型配置
// Small模型：256维，6层，4个专家，Top-2选择
model.printModelInfo();

// 前向传播
Variable output = model.forward(inputTokens);

// 获取负载均衡损失
float balancingLoss = model.computeTotalLoadBalancingLoss();

// 查看专家使用统计
model.printAllExpertStatistics();

// 高级配置
// Tiny模型：128维，4层，2专家，Top-1
// Medium模型：512维，8层，8专家，Top-2

5. Word2Vec词向量训练

使用Skip-gram模式训练词向量：

// 准备语料库
List<String> corpus = Arrays.asList(
    "机器", "学习", "是", "人工", "智能", "的", "重要", "分支",
    "深度", "学习", "是", "机器", "学习", "的", "子", "领域"
);

// 创建Word2Vec模型
Word2Vec word2vec = new Word2Vec(
    "word2vec_model",
    50,     // 词汇表大小
    10,     // 词向量维度
    Word2Vec.TrainingMode.SKIP_GRAM, // Skip-gram模式
    2,      // 上下文窗口大小
    true,   // 使用负采样
    5       // 负样本数量
);

// 构建词汇表和生成训练数据
word2vec.buildVocab(corpus);
List<Word2Vec.TrainingSample> samples = word2vec.generateTrainingSamples(corpus);

// 训练模型
Model model = new Model("word2vec_model", word2vec);
Optimizer optimizer = new SGD(model, 0.01f);
SoftmaxCrossEntropy lossFunc = new SoftmaxCrossEntropy();

// 简化训练循环
for (int epoch = 0; epoch < 100; epoch++) {
    for (Word2Vec.TrainingSample sample : samples) {
        Variable input = new Variable(new NdArray(new float[][]{{sample.input}}));
        Variable target = new Variable(new NdArray(new float[][]{{sample.target}}));

        Variable output = model.forward(input);
        Variable loss = lossFunc.loss(target, output);

        model.clearGrads();
        loss.backward();
        optimizer.update();
    }
}

// 查找相似词
List<String> similarWords = word2vec.mostSimilar("学习", 3);
System.out.println("与'学习'相似的词: " + similarWords);

6. GPT-2语言模型

创建小规模GPT-2模型：

// GPT-2参数设置
int vocabSize = 1000;   // 词汇表大小
int dModel = 128;       // 模型维度
int numLayers = 4;      // Transformer层数
int numHeads = 4;       // 注意力头数
int maxSeqLength = 64;  // 最大序列长度

// 创建GPT-2模型
GPT2Model gpt2 = new GPT2Model(
    "gpt2_small",
    vocabSize,
    dModel,
    numLayers,
    numHeads,
    maxSeqLength,
    0.1  // dropout率
);

// 模型初始化
gpt2.init();
gpt2.printModelInfo();

// 生成文本示例
NdArray inputTokens = new NdArray(new float[][]{{1, 2, 3, 4, 5}}); // token IDs
Variable logits = gpt2.forward(new Variable(inputTokens));

// 预测下一个token
int nextToken = gpt2.predictNextToken(inputTokens);
System.out.println("预测的下一个token ID: " + nextToken);

7. 强化学习示例

使用DQN解决CartPole问题：

// 创建环境
Environment env = new CartPoleEnvironment();

// 创建DQN智能体
DQNAgent agent = new DQNAgent(
    "CartPole_DQN",
    env.getStateDim(),      // 状态维度：4
    env.getActionDim(),     // 动作维度：2
    new int[]{128, 128},    // 隐藏层尺寸
    0.001f,                 // 学习率
    1.0f,                   // 初始探索率
    0.99f,                  // 折扣因子
    32,                     // 批次大小
    10000,                  // 缓冲区大小
    100                     // 目标网络更新频率
);

// 训练循环
for (int episode = 0; episode < 1000; episode++) {
    Variable state = env.reset();
    int totalReward = 0;

    while (!env.isDone()) {
        Variable action = agent.selectAction(state);
        Environment.StepResult result = env.step(action);

        Experience experience = new Experience(
            state, action, result.getReward(), 
            result.getNextState(), result.isDone()
        );

        agent.learn(experience);
        state = result.getNextState();
        totalReward += result.getReward().getValue().getNumber();
    }

    if (episode % 100 == 0) {
        System.out.println("Episode " + episode + ": Reward=" + totalReward + 
                          ", Epsilon=" + String.format("%.3f", agent.getCurrentEpsilon()));
    }
}

使用REINFORCE解决GridWorld问题：

// 创建GridWorld环境
Environment env = GridWorldEnvironment.createSimpleMaze();

// 创建REINFORCE智能体
REINFORCEAgent agent = new REINFORCEAgent(
    "GridWorld_REINFORCE",
    env.getStateDim(),      // 状态维度：2
    env.getActionDim(),     // 动作维度：4
    new int[]{64, 64},      // 隐藏层尺寸
    0.01f,                  // 学习率
    0.99f,                  // 折扣因子
    true                    // 使用基线
);

// 训练循环
for (int episode = 0; episode < 2000; episode++) {
    Variable state = env.reset();
    float episodeReward = 0;

    while (!env.isDone()) {
        Variable action = agent.selectAction(state);
        Environment.StepResult result = env.step(action);

        Experience experience = new Experience(
            state, action, result.getReward(), 
            result.getNextState(), result.isDone()
        );

        agent.learn(experience); // 存储经验
        state = result.getNextState();
        episodeReward += result.getReward().getValue().getNumber();
    }

    // REINFORCE在回合结束时学习
    agent.learnFromEpisode();

    if (episode % 200 == 0) {
        System.out.println("Episode " + episode + ": Reward=" + episodeReward);
    }
}

8. 模型序列化和管理

// 保存模型
model.saveModel("models/my_model.model");
model.saveModelCompressed("models/my_model_compressed.model");
model.saveParameters("models/parameters.params");

// 保存训练检查点
model.saveCheckpoint("checkpoints/epoch_100.ckpt", 100, 0.025f);

// 加载模型
Model loadedModel = Model.loadModel("models/my_model.model");

// 从检查点恢复
Model restoredModel = Model.resumeFromCheckpoint("checkpoints/epoch_100.ckpt");

// 模型信息管理
model.setDescription("这是一个图像分类模型");
model.updateTrainingInfo(100, 0.025f, "Adam", 0.001f);
model.addMetric("accuracy", 0.95f);
model.printModelInfo();

// 导出JSON报告
ModelInfoExporter.exportToJson(model, "reports/model_info.json");

9. CNN优化功能演示

使用增强的卷积神经网络：

```java // 创建带批量归一化的SimpleConvNet SimpleConvNet.ConvNetConfig config = new SimpleConvNet.ConvNetConfig() .filterNums(32, 64, 128) .useBatchNorm(true) .useResidual(true) .dropoutRate(0.3f);

SimpleConvNet convNet = SimpleConvNet.buildCustomConvNet( "enhanced_cnn", 3, 32, 32, 10, config);

// 使用深度可分离卷积 DepthwiseSeparableConvLayer dsConv = new DepthwiseSeparableConvLayer( "ds_conv", new Shape(16, 64, 32, 32), 128, 3, 1, 1);

// 批量归一化层 BatchNormLayer batchNorm = new BatchNormLayer("bn", 64, true); model.saveCheckpoint("checkpoints/epoch_100.ckpt", 100, 0.025f);

// 加载模型 Model loadedModel = Model.loadModel("model