# nox

**Repository Path**: zhengankun/nox

## Basic Information

- **Project Name**: nox
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-03-22
- **Last Updated**: 2026-04-07

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Nox - Symbolic Computation Graph Framework

一个支持符号形状、动态内存规划和 ONNX 导入的深度学习计算图框架。

## 目录

- [概述](#概述)
- [核心特性](#核心特性)
- [架构设计](#架构设计)
- [安装依赖](#安装依赖)
- [快速开始](#快速开始)
- [详细使用指南](#详细使用指南)
  - [1. 定义符号和域](#1-定义符号和域)
  - [2. 构建计算图](#2-构建计算图)
  - [3. 执行计算图](#3-执行计算图)
  - [4. ONNX 模型导入](#4-onnx-模型导入)
  - [5. 控制流支持](#5-控制流支持)
  - [6. 图可视化](#6-图可视化)
- [高级特性](#高级特性)
  - [内存管理机制](#内存管理机制)
  - [形状缓存优化](#形状缓存优化)
  - [状态张量](#状态张量)
  - [自动常量注册](#自动常量注册)
- [API 参考](#api-参考)
- [示例程序](#示例程序)
  - [comprehensive_example.py - 完整功能演示](#comprehensive_examplepy---完整功能演示)
  - [pytorch_cnn_onnx.py - ONNX 导入示例](#pytorch_cnn_onnxpy---onnx-导入示例)
  - [visualize_graph.py - 图可视化示例](#visualize_graphpy---图可视化示例)
  - [pytorch_if_example.py - PyTorch If 模型导出](#pytorch_if_examplepy---pytorch-if-模型导出)
  - [pytorch_loop_example.py - PyTorch Loop 模型导出](#pytorch_loop_examplepy---pytorch-loop-模型导出)
  - [pytorch_nested_control_flow.py - 嵌套控制流示例](#pytorch_nested_control_flowpy---嵌套控制流示例)
  - [pytorch_loop_patterns_example.py - 循环模式示例](#pytorch_loop_patterns_examplepy---循环模式示例)
  - [pytorch_break_continue_example.py - Break/Continue 模式示例](#pytorch_break_continue_examplepy---breakcontinue-模式示例)
- [项目结构](#项目结构)
- [性能优化建议](#性能优化建议)
- [常见问题](#常见问题)
- [未来计划](#未来计划)

## 概述

这是一个专为深度学习模型设计的符号计算图框架，支持动态形状、符号化维度、内存优化和控制流操作。框架的核心思想是：

1. **符号形状**：张量形状可以是包含 `sympy.Symbol` 的表达式，支持动态维度
2. **域约束**：通过 `Domain` 为符号定义取值范围，用于内存预分配
3. **内存规划**：根据符号的最大可能值一次性分配内存，避免运行时动态分配
4. **控制流支持**：原生支持 `If` 和 `Loop` 等控制流操作
5. **ONNX 兼容**：可导入 ONNX 模型，自动处理动态维度
6. **图可视化**：支持导出 Graphviz DOT 格式，直观展示计算图结构

该框架特别适合以下场景：
- 需要处理动态形状的模型（如变长序列、动态 batch）
- 嵌入式或移动端部署，需要可预测的内存占用
- 需要精细控制内存布局的性能敏感应用
- 需要将 PyTorch 模型转换为静态内存分配图的场景

## 核心特性

### 1. 符号化形状系统

- 张量形状可以使用 `sympy.Symbol` 作为维度
- 支持算术表达式，如 `(N, C, H//2, W//2)`
- 运行时根据实际输入解析具体形状

```python
from sympy import Symbol

N, C, H, W = Symbol('N'), Symbol('C'), Symbol('H'), Symbol('W')
input_shape = (N, C, H, W)
conv_out_shape = (N, 64, H//2, W//2)
```

### 2. 域约束系统

通过 `Domain` 为符号定义取值范围：

```python
from nox import Domain

domain = Domain({
    N: (1, 32),      # 区间范围 1~32
    C: 3,            # 固定值 3
    H: (224, 224),   # 固定 224
    W: [224, 256]    # 离散值 224 或 256
})
```

### 3. 自动内存规划

- 使用 `BestFitAllocator` 进行内存分配
- 根据域的最大值预分配内存
- 支持设备（CPU/CUDA）分离的内存管理
- 自动处理子图内存复用

### 4. 控制流算子

- **If 操作**：条件分支，支持 `then` 和 `else` 分支
- **Loop 操作**：循环操作，支持固定迭代次数和条件终止
- 子图独立内存规划，自动复用
- 支持嵌套控制流（Loop 内嵌套 If，If 内嵌套 Loop 等）

### 5. ONNX 模型导入

- 自动解析 ONNX 模型中的动态维度
- 支持子图（If、Loop 节点）
- 保留符号信息用于内存规划
- 支持嵌套控制流模型导入
- 自动验证 ONNX 控制流结构正确性

#### ONNX 控制流验证

导入 ONNX 模型后，框架会自动验证控制流结构的正确性：

```python
from nox.onnx import load_onnx

graph = load_onnx("model.onnx", device='cuda')
# 验证结果会自动记录在日志中
# 也可以手动调用: result = graph.validate_onnx_control_flow()
```

验证规则包括：
- **If 算子**：分支输出数量、类型必须匹配
- **Loop 算子**：carry_in/carry_out 类型必须一致，cond_out 为 BOOL 类型
- **嵌套控制流**：递归验证所有子图

### 6. 图可视化

- 支持导出 Graphviz DOT 格式
- 支持嵌套子图的层级显示
- 自动标注算子类型、输入输出和形状信息

```python
from nox import visualize_graph

# 导出为 DOT 文件
visualize_graph(graph, filename="my_graph", format="dot")

# 或导出为 SVG 图片
visualize_graph(graph, filename="my_graph", format="svg", view=True)
```

### 7. 结构化日志

- 基于 `structlog` 的结构化日志系统
- 自动添加时间戳、日志级别、调用位置信息
- 彩色控制台输出，便于调试

```python
from nox.utils.logging import get_logger

logger = get_logger(__name__)
logger.debug("Processing tensor", name=tensor.name, shape=tensor.shape)
```

## 架构设计

### 核心类层次结构

```
Graph (计算图)
  ├── Tensor (张量)
  │    ├── shape (形状，可包含 sympy.Symbol)
  │    ├── dtype (数据类型)
  │    ├── device (设备)
  │    └── offset (内存偏移)
  ├── Op (算子)
  │    ├── Conv2dOp
  │    ├── ReLUOp
  │    ├── AddOp
  │    ├── IfOp
  │    ├── LoopOp
  │    └── ...
  └── Domain (符号域)

MemoryManager (内存管理器)
  ├── GlobalMemoryManager (全局单例)
  ├── SubMemoryManager (子图内存管理器)
  └── BestFitAllocator (最佳适配分配器)
```

### 工作流程

1. **图构建阶段**
   - 创建 `Graph` 对象，传入 `domain` 参数
   - 添加算子和张量
   - 自动进行形状推导

2. **内存规划阶段**（在 `execute_graph` 内部自动执行）
   - 根据 `Domain` 计算每个张量的最大尺寸
   - 使用最佳适配算法分配偏移
   - 递归处理子图内存
   - 创建实际的设备缓冲区

3. **解析执行阶段**
   - 根据实际输入解析符号映射
   - 设置输入数据
   - 执行计算图
   - 返回输出结果字典

## 安装依赖

### 基础依赖

```bash
pip install torch sympy
```

### 可选依赖（ONNX 支持）

```bash
pip install onnx
```

### 可选依赖（图可视化）

```bash
pip install graphviz
```

### 完整安装

```bash
pip install torch sympy onnx graphviz
```

## 快速开始

### 示例 1：简单卷积网络

```python
import torch
from sympy import Symbol
from nox import Domain, DataType, Tensor, Graph, conv2d, relu, execute_graph

# 1. 定义符号和域
N, C, H, W = Symbol('N'), Symbol('C'), Symbol('H'), Symbol('W')
domain = Domain({
    N: (1, 4),      # batch size 1-4
    C: 3,           # 固定 3 通道
    H: 224,         # 固定高度
    W: 224          # 固定宽度
})

# 2. 构建图（domain 会自动触发内存规划）
with Graph("simple_cnn", domain=domain) as graph:
    # 输入张量
    input_t = Tensor("input", shape=(N, C, H, W), dtype=DataType.FLOAT32, device='cuda')
    weight = Tensor("weight", shape=(64, C, 3, 3), dtype=DataType.FLOAT32, device='cuda')
    bias = Tensor("bias", shape=(64,), dtype=DataType.FLOAT32, device='cuda')

    # 计算
    x = conv2d(input_t, weight, bias, stride=2, padding=1, name="conv1")
    x = relu(x, name="relu1")

    graph.add_output(x)

# 3. 执行（自动解析形状和设置输入）
input_val = torch.randn(1, 3, 224, 224, device='cuda')
weight_val = torch.randn(64, 3, 3, 3, device='cuda')
bias_val = torch.randn(64, device='cuda')

outputs = execute_graph(graph, input_tensors={
    "input": input_val,
    "weight": weight_val,
    "bias": bias_val,
})

# 4. 获取结果
output = outputs["conv1_out"]
print(f"Output shape: {output.shape}")
```

### 示例 2：带控制流的网络

```python
from nox import If

with Graph("control_flow", domain=domain) as graph:
    input_t = Tensor("input", shape=(N, C, H, W), dtype=DataType.FLOAT32, device='cuda')
    cond_t = Tensor("cond", shape=(1,), dtype=DataType.INT32, device='cuda')

    x = conv2d(input_t, weight1, bias1, stride=2, padding=1, name="conv1")

    # If 条件分支
    if_ctx = If(cond_t, "branch")
    with if_ctx:
        with if_ctx.then() as then_graph:
            y = conv2d(x, weight2, bias2, stride=1, padding=1, name="conv2_then")
            then_graph.add_output(y)

        with if_ctx.else_() as else_graph:
            else_graph.add_output(x)

    graph.add_output(if_ctx.output)

# 执行
outputs = execute_graph(graph, input_tensors={
    "input": input_val,
    "cond": torch.tensor([1], dtype=torch.int32, device='cuda'),
    "weight1": weight1_val,
    "bias1": bias1_val,
    "weight2": weight2_val,
    "bias2": bias2_val,
})
```

### 示例 3：图可视化

```python
from nox import visualize_graph

# 导出为 DOT 文件
visualize_graph(graph, filename="my_graph", format="dot")

# 或直接导出为 SVG 并打开
visualize_graph(graph, filename="my_graph", format="svg", view=True)
```

生成的 DOT 文件支持：
- 嵌套子图的层级显示
- 算子类型、输入输出标注
- 形状和数据类型信息

```python
from sympy import Symbol

# 创建符号
batch = Symbol('batch')
channels = Symbol('channels')
height = Symbol('height')
width = Symbol('width')

# 组合成形状
input_shape = (batch, channels, height, width)
```

#### 域定义

```python
# 方式1：区间范围
domain = Domain({
    batch: (1, 32),      # 1 到 32
    channels: (3, 3),    # 固定为 3
    height: (64, 512),   # 64 到 512
    width: (64, 512)
})

# 方式2：离散值列表
domain = Domain({
    batch: [1, 2, 4, 8],   # 只能是这些值
    height: [224, 256, 288]
})

# 方式3：混合方式
domain = Domain({
    batch: (1, 8),
    height: [224, 256],
    width: (224, 512)
})

# 方式4：带约束
from sympy import Ge, Le

domain = Domain(
    {batch: (1, 32), height: (64, 512)},
    constraints=[
        height >= batch * 32,     # 高度至少是 batch 的 32 倍
        height <= batch * 64      # 高度最多是 batch 的 64 倍
    ]
)
```

### 2. 构建计算图

#### 创建图

```python
# 创建主图
with Graph("main", domain=domain) as graph:
    # 添加输入
    input_t = Tensor("input", shape=(N, C, H, W), device='cuda')
    
    # 添加常量
    weight = Tensor("weight", shape=(64, C, 3, 3), device='cuda')
    graph.add_const(weight, torch.randn(64, 3, 3, 3))
    
    # 添加算子
    x = conv2d(input_t, weight, stride=2, padding=1)
    x = relu(x)
    
    # 设置输出
    graph.add_output(x)
```

#### 算子工厂函数

框架提供了以下算子工厂函数：

| 算子 | 函数 | 描述 |
|------|------|------|
| 卷积 | `conv2d(input, weight, bias, stride, padding, dilation, groups)` | 2D 卷积 |
| ReLU | `relu(input)` | ReLU 激活 |
| 加法 | `add(a, b)` | 张量加法 |
| 比较 | `greater(a, b)` | 大于比较 |
| 拷贝 | `copy_to(input, target_device)` | 跨设备拷贝 |
| NMS | `nms(boxes, scores, iou_threshold)` | 非极大值抑制 |
| 状态读写 | `set_state(key, value)`, `get_state(key)` | 全局状态访问 |

### 3. 内存规划

内存规划是框架的核心功能，它根据符号的最大可能值预先分配内存。

```python
# 规划内存（使用 domain 中的最大值）
buffers, offsets = graph.plan_memory()
GlobalMemoryManager.initialize(buffers, offsets)

# 打印内存使用情况
print(f"Total GPU memory: {buffers['cuda'].numel()} bytes")
print(f"Total CPU memory: {buffers['cpu'].numel()} bytes")
```

内存规划过程：

1. **计算最大尺寸**：遍历所有张量，使用 `domain.max_point()` 计算每个维度的最大值
2. **分配偏移**：使用最佳适配算法为每个张量分配内存偏移
3. **处理子图**：递归处理子图的内存，并为其分配 workspace
4. **创建设备缓冲区**：为每个设备创建 `torch.uint8` 类型的连续内存缓冲区

### 4. 图执行

#### 符号解析

在执行前，需要根据实际输入解析符号映射：

```python
# 准备实际输入形状
input_shapes = {
    input_t.id: (2, 3, 224, 224),  # batch=2
    # ... 其他输入
}

# 解析符号
graph.resolve_symbols(input_shapes)

# 现在所有张量都有 concrete_shape 和实际大小
for t in graph.tensors:
    if t.concrete_shape:
        print(f"{t.name}: {t.concrete_shape}, size={t.size}")
```

#### 数据设置和获取

```python
# 设置输入数据
GlobalMemoryManager.set_tensor(input_t, input_data)

# 获取输出数据
output_data = GlobalMemoryManager.get_tensor(output_t)
```

#### 执行图

```python
# 执行整个图
execute_graph(graph)
```

### 5. ONNX 模型导入

#### 基本导入

```python
# 导入 ONNX 模型
onnx_graph = load_onnx("model.onnx", device='cuda', domain=domain)
```

#### 处理动态维度

ONNX 模型中的动态维度会自动转换为符号：

```python
# ONNX 模型定义
# input: ["batch", 3, 224, 224]  # batch 是动态的

# 加载时指定 batch 的范围
batch = Symbol('batch')
domain = Domain({batch: (1, 16)})
graph = load_onnx("model.onnx", device='cuda', domain=domain)
```

#### 支持的控制流节点

- **If 节点**：条件分支，自动转换为 `IfOp`
- **Loop 节点**：循环操作，自动转换为 `LoopOp`

### 6. 控制流支持

#### If 操作

```python
if_ctx = If(condition, name="my_if")
with if_ctx:
    # then 分支
    with if_ctx.then() as then_graph:
        # 在 then 分支中添加算子
        y = conv2d(x, weight)
        then_graph.add_output(y)
    
    # else 分支
    with if_ctx.else_() as else_graph:
        else_graph.add_output(x)

# 获取输出
result = if_ctx.output
```

注意事项：
- then 和 else 分支必须输出相同形状和类型的张量
- 分支内部可以包含任意复杂的计算图
- 分支可以访问外部张量（自动捕获）

#### Loop 操作

```python
# 固定迭代次数的循环
loop_ctx = Loop(max_iter=10, name="my_loop")
loop_ctx.init_inputs = [state1, state2]  # 初始状态

with loop_ctx:
    with loop_ctx.body() as body_graph:
        # 获取当前状态
        s1 = body_graph.inputs[0]  # 或使用 Tensor 对象
        
        # 循环体计算
        new_s1 = add(s1, one)
        
        # 输出新状态
        body_graph.add_output(new_s1)

# 获取最终状态
final_state = loop_ctx.outputs[0]
```

循环特性：
- 支持多状态变量
- 自动处理状态更新
- 支持条件终止（通过 `max_iter=None` 并使用输入作为迭代计数）

## 高级特性

### 内存管理机制

#### BestFitAllocator

使用最佳适配算法管理内存块：

```python
alloc = BestFitAllocator()

# 分配内存
offset1 = alloc.allocate(1024)   # 分配 1024 字节
offset2 = alloc.allocate(2048)   # 分配 2048 字节

# 释放内存
alloc.free(offset1, 1024)

# 重新分配（复用已释放的空间）
offset3 = alloc.allocate(512)    # 可能会复用 offset1 的位置
```

#### 内存复用策略

1. **生命周期分析**：计算每个张量的最后使用算子
2. **空间复用**：不同生命周期的张量可以共享内存
3. **子图复用**：控制流分支共享 workspace 内存

### 子图内存复用

控制流算子（If、Loop）的子图共享同一个 workspace：

```python
# 子图独立规划内存
then_graph.plan_memory()   # 返回子图所需内存大小
else_graph.plan_memory()

# 父图分配最大子图所需内存
max_sub_mem = max(then_mem, else_mem)
sub_workspace = parent_alloc.allocate(max_sub_mem)

# 子图共享同一个 workspace
then_graph.use_workspace(sub_workspace)
else_graph.use_workspace(sub_workspace)
```

### 状态张量

状态张量用于在控制流之间传递持久化状态：

```python
# 创建状态张量
state = Tensor("counter", shape=(1,), dtype=DataType.FLOAT32, device='cuda')
state_tensors = {"counter": state}

with Graph("main", state_tensors=state_tensors) as graph:
    # 读取状态
    counter = get_state("counter")
    
    # 更新状态
    new_counter = add(counter, one)
    set_state("counter", new_counter)
```

状态张量的特点：
- 在整个图生命周期内持久化
- 可在子图中访问和修改
- 自动处理内存分配

## API 参考

### Graph 类

| 方法 | 描述 |
|------|------|
| `__init__(name, parent, state_tensors, domain)` | 创建图 |
| `add_input(tensor)` | 添加输入张量 |
| `add_output(tensor)` | 添加输出张量 |
| `add_const(tensor, value)` | 添加常量张量 |
| `add_op(op)` | 添加算子 |
| `plan_memory()` | 规划内存，返回 (buffers, offsets) |
| `resolve_symbols(input_shapes)` | 根据输入解析符号 |

### Tensor 类

| 属性 | 描述 |
|------|------|
| `name` | 张量名称 |
| `symbolic_shape` | 符号形状 |
| `concrete_shape` | 具体形状（解析后） |
| `dtype` | 数据类型 |
| `device` | 设备 |
| `offset` | 内存偏移 |
| `size` | 实际大小 |
| `max_size` | 最大大小 |

### Domain 类

| 方法 | 描述 |
|------|------|
| `__init__(ranges, constraints)` | 创建域 |
| `enumerate_points()` | 枚举所有有效点 |
| `max_point()` | 获取最大值点 |
| `volume()` | 计算有效点数量 |

### MemoryManager 类

| 方法 | 描述 |
|------|------|
| `get_tensor(tensor)` | 获取张量数据 |
| `set_tensor(tensor, value)` | 设置张量数据 |
| `share_tensor(tensor, value)` | 共享张量数据 |

## 示例程序

### 完整示例：ResNet-50 风格的网络

```python
import torch
from sympy import Symbol
import framework as fw

# 定义符号
N, C, H, W = Symbol('N'), Symbol('C'), Symbol('H'), Symbol('W')
domain = fw.Domain({
    N: (1, 32),
    C: 3,
    H: 224,
    W: 224
})

with fw.Graph("resnet_style", domain=domain) as graph:
    # 输入
    input_t = fw.Tensor("input", shape=(N, C, H, W), device='cuda')
    
    # 第一个卷积块
    w1 = fw.Tensor("conv1_weight", shape=(64, C, 7, 7), device='cuda')
    b1 = fw.Tensor("conv1_bias", shape=(64,), device='cuda')
    x = fw.conv2d(input_t, w1, b1, stride=2, padding=3)
    x = fw.relu(x)
    
    # 残差块
    w2 = fw.Tensor("conv2_weight", shape=(64, 64, 3, 3), device='cuda')
    b2 = fw.Tensor("conv2_bias", shape=(64,), device='cuda')
    residual = x
    x = fw.conv2d(x, w2, b2, padding=1)
    x = fw.relu(x)
    x = fw.add(x, residual)
    
    # 全局池化
    # ... 更多操作
    
    graph.add_output(x)

# 规划内存
buffers, offsets = graph.plan_memory()
fw.GlobalMemoryManager.initialize(buffers, offsets)

# 执行
input_data = torch.randn(1, 3, 224, 224, device='cuda')
fw.GlobalMemoryManager.set_tensor(input_t, input_data)
# ... 设置其他数据

fw.execute_graph(graph)
```

### ONNX 模型导入示例

```python
import framework as fw
from sympy import Symbol

# 假设有一个 ONNX 模型，输入形状为 [batch, 3, 224, 224]
batch = Symbol('batch')
domain = fw.Domain({batch: (1, 16)})

# 导入模型
graph = fw.load_onnx("resnet50.onnx", device='cuda', domain=domain)

# 内存规划
buffers, offsets = graph.plan_memory()
fw.GlobalMemoryManager.initialize(buffers, offsets)

# 执行
input_shapes = {graph.input_ids[0]: (8, 3, 224, 224)}  # batch=8
graph.resolve_symbols(input_shapes)

# 设置输入数据并执行
# ...
```

## 性能优化建议

### 1. 合理设置域范围

- 域范围越小，内存占用越少
- 使用离散值而不是大区间，提高内存利用率
- 利用约束表达式减少无效形状组合

```python
# 不好的做法
domain = Domain({N: (1, 1000)})  # 可能浪费大量内存

# 好的做法
domain = Domain({N: [1, 2, 4, 8, 16, 32]})  # 只考虑实际使用的 batch size
```

### 2. 使用内存复用

- 手动管理张量生命周期
- 尽量复用张量而不是创建新的

### 3. 控制流优化

- 避免深层嵌套的控制流
- 尽量使用固定迭代次数的循环而不是条件循环

### 4. 设备选择

- 将常量和权重放在计算设备（GPU）
- 将不需要加速的数据放在 CPU

## 常见问题

### Q1: 如何调试符号解析问题？

```python
# 打印符号映射
print(graph.symbol_map)

# 打印张量的符号和具体形状
for t in graph.tensors:
    print(f"{t.name}: symbolic={t.symbolic_shape}, concrete={t.concrete_shape}")
```

### Q2: 内存规划失败，提示 "Domain not set for graph"

确保在创建 Graph 时传递了 domain 参数：

```python
graph = Graph("my_graph", domain=domain)
```

### Q3: ONNX 导入时出现 "Unsupported op type"

扩展 `ONNXLoader.op_map` 添加对新算子的支持：

```python
class MyONNXLoader(ONNXLoader):
    def __init__(self, model_path, device, domain):
        super().__init__(model_path, device, domain)
        self.op_map['MyOp'] = self._convert_myop
    
    def _convert_myop(self, node):
        # 实现转换逻辑
        pass
```

### Q4: 如何处理变长序列（如 LSTM）？

使用符号表示序列长度：

```python
seq_len = Symbol('seq_len')
domain = Domain({seq_len: (1, 512)})
input_shape = (batch, seq_len, hidden_size)
```

### Q5: 如何实现自定义算子？

继承 `Op` 类并实现必要方法：

```python
from nox.ops.base import Op

class MyCustomOp(Op):
    def infer_shape(self):
        # 推导输出形状
        self.outputs[0].set_shape(self.inputs[0].shape)

    def infer_dtype(self):
        self.outputs[0].set_dtype(self.inputs[0].dtype)

    def execute_(self, input_tensors, output_tensors):
        # 执行计算
        result = input_tensors[0] * 2
        output_tensors[0].copy_(result)
```

## 项目结构

```
nox/
├── __init__.py          # 主包入口，导出公共 API
├── core/                # 核心类型和工具
│   ├── __init__.py
│   ├── tensor.py        # Tensor 类定义
│   ├── dtype.py         # DataType 枚举
│   ├── domain.py        # Domain 类和 expr_min_max 函数
│   ├── utils.py         # get_size 等工具函数
│   └── context.py       # 图上下文管理
├── graph/
│   ├── __init__.py
│   └── graph.py         # Graph 类定义
├── memory/
│   ├── __init__.py
│   ├── allocator.py     # BestFitAllocator 内存分配器
│   └── manager.py       # MemoryManager 和 GlobalMemoryManager
├── execution/
│   ├── __init__.py
│   └── execute.py       # execute_graph 函数
├── ops/
│   ├── __init__.py
│   ├── base.py          # Op 基类
│   ├── nn.py            # 神经网络算子 (Conv2dOp, ReLUOp)
│   ├── elementwise.py   # 元素级算子 (AddOp, GreaterOp)
│   ├── control_flow.py  # 控制流算子 (IfOp, LoopOp)
│   └── ...
├── api/
│   ├── __init__.py
│   ├── factory.py       # 算子工厂函数 (conv2d, relu 等)
│   └── control_flow.py  # If, Loop 上下文管理器
├── onnx/
│   ├── __init__.py
│   └── loader.py        # ONNX 模型加载器
├── visualization/
│   ├── __init__.py
│   └── visualizer.py    # GraphVisualizer 图可视化
├── utils/
│   └── debug.py         # 调试工具
└── examples/            # 示例程序
    ├── main.py          # 主示例（完整功能演示）
    ├── onnx_conversion.py  # ONNX 导入示例
    └── visualize_graph.py  # 图可视化示例
```

### 核心模块说明

| 模块 | 说明 |
|------|------|
| `nox.core` | 核心类型定义（Tensor, DataType, Domain） |
| `nox.graph` | Graph 类，计算图容器 |
| `nox.memory` | 内存管理（分配器、管理器） |
| `nox.execution` | 图执行逻辑 |
| `nox.ops` | 算子定义和实现 |
| `nox.api` | 用户友好的 API 接口 |
| `nox.onnx` | ONNX 模型导入 |
| `nox.visualization` | Graphviz 图可视化 |

## 示例程序

### comprehensive_example.py - 完整功能演示

演示完整的 Nox 工作流程，包括：
- 符号形状定义
- 控制流（If 和 Loop）
- 状态张量
- 多 batch size 测试

```bash
python examples/comprehensive_example.py
```

### pytorch_cnn_onnx.py - ONNX 导入示例

演示从 PyTorch 导出 ONNX 模型并导入 Nox：
- 动态轴处理
- 多 batch size 验证
- 与 PyTorch 结果对比

```bash
python examples/pytorch_cnn_onnx.py
```

### visualize_graph.py - 图可视化示例

演示如何将计算图导出为 Graphviz DOT/SVG 格式：
- 简单 CNN 可视化
- 嵌套控制流可视化（If 内嵌套 Loop）

```bash
python examples/visualize_graph.py
# 输出：dumps/visualization/simple_cnn.dot
#       dumps/visualization/with_control_flow.dot
```

### pytorch_if_example.py - PyTorch If 模型导出

演示包含条件分支的 PyTorch 模型导出：
- 使用 `torch.jit.script` 保留控制流
- 导出为 ONNX If 算子
- 在 Nox 中执行并验证

```bash
python examples/pytorch_if_example.py
```

### pytorch_loop_example.py - PyTorch Loop 模型导出

演示包含动态循环的 PyTorch 模型导出：
- 使用 while 循环生成 ONNX Loop 算子
- 可变迭代次数
- 在 Nox 中执行并验证

```bash
python examples/pytorch_loop_example.py
```

### pytorch_nested_control_flow.py - 嵌套控制流示例

演示包含嵌套控制流的 PyTorch 模型导出：
- 三层嵌套结构（Loop 内嵌套 If，If 内嵌套 Loop）
- 使用 `torch.jit.script` 保留控制流
- 导出为 ONNX 并在 Nox 中执行验证
- 验证嵌套控制流的正确性

```bash
python examples/pytorch_nested_control_flow.py
```

### pytorch_loop_patterns_example.py - 循环模式示例

演示多种循环模式的 PyTorch 模型导出：
- 简单累加循环
- 带条件累加的循环（Loop 内嵌套 If）
- 嵌套循环
- 自动验证所有模式在 Nox 中的正确性

```bash
python examples/pytorch_loop_patterns_example.py
```

## 项目结构

```
nox/
├── __init__.py          # 主包入口，导出公共 API
├── core/                # 核心类型和工具
│   ├── __init__.py
│   ├── tensor.py        # Tensor 类定义
│   ├── dtype.py         # DataType 枚举
│   ├── domain.py        # Domain 类和 expr_min_max 函数
│   ├── utils.py         # get_size 等工具函数
│   └── context.py       # 图上下文管理
├── graph/
│   ├── __init__.py
│   └── graph.py         # Graph 类定义
├── memory/
│   ├── __init__.py
│   ├── allocator.py     # BestFitAllocator 内存分配器
│   └── manager.py       # MemoryManager 和 GlobalMemoryManager
├── execution/
│   ├── __init__.py
│   └── execute.py       # execute_graph 函数
├── ops/
│   ├── __init__.py
│   ├── base.py          # Op 基类
│   ├── nn.py            # 神经网络算子 (Conv2dOp, ReLUOp)
│   ├── elementwise.py   # 元素级算子 (AddOp, GreaterOp)
│   ├── control_flow.py  # 控制流算子 (IfOp, LoopOp)
│   └── ...
├── api/
│   ├── __init__.py
│   ├── factory.py       # 算子工厂函数 (conv2d, relu 等)
│   └── control_flow.py  # If, Loop 上下文管理器
├── onnx/
│   ├── __init__.py
│   └── loader.py        # ONNX 模型加载器
├── visualization/
│   ├── __init__.py
│   └── visualizer.py    # GraphVisualizer 图可视化
├── utils/
│   └── debug.py         # 调试工具
└── examples/            # 示例程序
    ├── comprehensive_example.py          # 完整功能演示
    ├── pytorch_cnn_onnx.py               # ONNX 导入示例
    ├── visualize_graph.py                # 图可视化示例
    ├── pytorch_if_example.py             # PyTorch If 模型导出
    ├── pytorch_loop_example.py           # PyTorch Loop 模型导出
    ├── pytorch_nested_control_flow.py    # 嵌套控制流示例
    └── pytorch_loop_patterns_example.py  # 循环模式示例
```

## 未来计划

### 短期计划

- [ ] 支持更多 ONNX 算子（BatchNorm、Pooling、MatMul 等）
- [ ] 添加图优化 Pass（常量折叠、算子融合）
- [ ] 支持更复杂的内存复用策略
- [ ] 添加性能分析和可视化工具

### 中期计划

- [ ] 支持多设备自动并行
- [ ] 添加量化支持（INT8、INT4）
- [ ] 实现动态形状的运行时优化
- [ ] 支持 TensorRT 后端

### 长期计划

- [ ] 完整的深度学习框架集成
- [ ] 支持训练和自动微分
- [ ] 分布式训练支持
- [ ] 边缘设备部署优化

## 贡献指南

我们欢迎任何形式的贡献！

### 如何贡献

1. Fork 项目
2. 创建功能分支 (`git checkout -b feature/AmazingFeature`)
3. 提交更改 (`git commit -m 'Add some AmazingFeature'`)
4. 推送到分支 (`git push origin feature/AmazingFeature`)
5. 创建 Pull Request

### 代码规范

- 遵循 PEP 8 规范
- 添加必要的注释和文档字符串
- 编写单元测试
- 更新相关文档

### 报告问题

请在 GitHub Issues 中报告问题，包含：
- 问题描述
- 复现步骤
- 预期行为
- 实际行为
- 环境信息（Python 版本、PyTorch 版本等）

## 许可证

本项目采用 MIT 许可证。

## 致谢

- PyTorch 团队提供的深度学习框架
- SymPy 团队提供的符号计算库
- ONNX 社区提供的模型交换格式

## 联系方式

如有问题或建议，请通过以下方式联系我们：

- GitHub Issues: [项目 Issues 页面]
- Email: [项目邮箱]

---

**注意**：本文档中的代码示例可能需要根据实际环境进行调整。建议在实际使用前运行示例程序验证环境配置。
```

这个 README 提供了全面的文档，涵盖了框架的核心概念、使用方法、API 参考和最佳实践，帮助用户快速上手并深入理解框架的设计思想。