graph2net 0.1.0

从图结构到PyTorch模型的编译库 - 能连上就是合法的

Graph2Net

从图结构到PyTorch模型的编译库

设计理念

"能连上就是合法的" - Graph2Net 自动处理维度不匹配，让你专注于模型架构设计而非繁琐的维度计算。

核心特性

• 图结构定义模型: 使用JSON定义模型架构，节点是PyTorch模块，边代表张量流向
• 自动维度适配: 自动推断维度，智能插入适配器（Reshape/View/线性投影）
• 全局配置系统: 自动选择最优的归一化类型和Dropout放置策略
• 子图复用: 支持子图定义和复用，参数可覆盖
• 现代架构支持: Transformer (MoE, RoPE, Flash Attention), Diffusion (UNet), CNN, RNN
• 编译式执行: 图JSON编译为PyTorch模型实例，支持静态优化

快速开始

安装

pip install graph2net

基础示例

import torch
from graph2net import compile_graph

# 定义模型结构
graph = {
    "nodes": {
        "input": {"type": "Input", "params": {"shape": [2, 784]}},
        "fc1": {"type": "Linear", "params": {"out_features": 256, "activation": "relu"}},
        "fc2": {"type": "Linear", "params": {"out_features": 128, "activation": "relu"}},
        "output": {"type": "Linear", "params": {"out_features": 10, "activation": "softmax"}}
    },
    "edges": [
        {"from": "input", "to": "fc1"},
        {"from": "fc1", "to": "fc2"},
        {"from": "fc2", "to": "output"}
    ]
}

# 编译为PyTorch模型
model = compile_graph(graph)

# 使用模型
x = torch.randn(2, 784)
output = model(x)
print(output.shape)  # torch.Size([2, 10])

全局配置

from graph2net import configure

# 配置为Transformer网络（自动使用RMSNorm等）
configure(network_type="transformer", dropout_rate=0.1)

# 节点使用全局配置
graph = {
    "nodes": {
        "input": {"type": "Input", "params": {"shape": [2, 512]}},
        # norm="auto" 自动选择归一化类型，dropout=None 使用全局配置
        "fc": {"type": "Linear", "params": {"out_features": 512, "norm": "auto", "dropout": None}}
    },
    "edges": [{"from": "input", "to": "fc"}]
}

子图复用

# 定义可复用的子图
graph = {
    "subgraphs": {
        "resblock": {
            "nodes": {
                "input": {"type": "Input", "params": {"shape": [null, 64]}},
                "fc1": {"type": "Linear", "params": {"out_features": 64, "activation": "relu"}},
                "fc2": {"type": "Linear", "params": {"out_features": 64}},
                "output": {"type": "Output", "params": {}}
            },
            "edges": [
                {"from": "input", "to": "fc1"},
                {"from": "fc1", "to": "fc2"},
                {"from": "fc2", "to": "output"}
            ]
        }
    },
    "nodes": {
        "input": {"type": "Input", "params": {"shape": [2, 64]}},
        "res1": {"type": "resblock", "params": {}},  # 使用子图
        "res2": {"type": "resblock", "params": {}},  # 复用子图
        "output": {"type": "Output", "params": {}}
    },
    "edges": [
        {"from": "input", "to": "res1"},
        {"from": "res1", "to": "res2"},
        {"from": "res2", "to": "output"}
    ]
}

支持的节点类型

基础层

• Linear - 全连接层，支持激活、归一化、Dropout
• Conv - 卷积层（自动推断1D/2D/3D）
• Flatten, Reshape, Identity

归一化

• BatchNorm, LayerNorm, GroupNorm, InstanceNorm, RMSNorm

注意力

• SelfAttention, CrossAttention
• TransformerEncoder, TransformerDecoder
• MultiHeadAttention (支持Flash Attention)

RNN

• RNN, LSTM, GRU (统一接口)

Diffusion

• UNet - UNet架构
• ResBlock - 残差块
• TimestepEmbedding - 时间步嵌入
• DDPMSampler, DDIMSampler - 采样器

合并

• 自动多输入合并（Add/Concat/Multiply）

网络类型配置

网络类型	归一化	Dropout放置
CNN	BatchNorm	激活后
Transformer	RMSNorm	归一化后
Diffusion	GroupNorm	归一化前
MLP	BatchNorm	激活后
RNN	LayerNorm	归一化后

自动维度适配

Graph2Net 自动处理以下维度不匹配情况：

• 形状不匹配: 自动插入Reshape/View
• 特征维度不匹配: 自动插入线性投影
• Conv ↔ Transformer: 自动Flatten + 投影
• 多输入合并: 自动适配后Add/Concat

测试

# 运行所有测试
python test_graph2net.py

# 或使用pytest
pytest tests/

许可证

MIT License - 详见 LICENSE 文件

贡献

欢迎提交Issue和Pull Request！

致谢

感谢所有为Graph2Net做出贡献的开发者！