Torch-Struct：结构化预测库

一个为深度学习应用提供核心结构化预测算法的、经过充分测试且支持 GPU 的实现库。

• HMM / 线性链 CRF
• HSMM / 半马尔可夫 CRF
• 依存句法树 CRF
• PCFG 二叉树 CRF
• ...

设计用于作为高效批处理层集成到其他 PyTorch 代码中。

描述该方法论的教程论文。

快速入门

!pip install -qU git+https://github.com/harvardnlp/pytorch-struct
# 可选的 CUDA 内核，用于 FastLogSemiring
!pip install -qU git+https://github.com/harvardnlp/genbmm
# 用于绘图。
!pip install -q matplotlib

import torch
from torch_struct import DependencyCRF, LinearChainCRF
import matplotlib.pyplot as plt
def show(x): plt.imshow(x.detach())

# 生成一些数据。
vals = torch.zeros(2, 10, 10) + 1e-5
vals[:, :5, :5] = torch.rand(5)
vals[:, 5:, 5:] = torch.rand(5) 
dist = DependencyCRF(vals.log())
show(dist.log_potentials[0])

# 计算边缘概率
show(dist.marginals[0])

# 计算最大值
show(dist.argmax.detach()[0])

# 计算评分和枚举（前向/内部）
log_partition = dist.partition
max_score = dist.log_prob(dist.argmax)

# 计算样本 
show(dist.sample((1,)).detach()[0, 0])

# 库内置了填充/掩码功能。
dist = DependencyCRF(vals, lengths=torch.tensor([10, 7]))
show(dist.marginals[0])
plt.show()
show(dist.marginals[1])

# 其他多种结构化预测方法
chain = torch.zeros(2, 10, 10, 10) + 1e-5
chain[:, :, :, :] = vals.unsqueeze(-1).exp()
chain[:, :, :, :] += torch.eye(10, 10).view(1, 1, 10, 10) 
chain[:, 0, :, 0] = 1
chain[:, -1,9, :] = 1
chain = chain.log()

dist = LinearChainCRF(chain)
show(dist.marginals.detach()[0].sum(-1))

库

完整文档：http://nlp.seas.harvard.edu/pytorch-struct/

当前实现的分布包括：

• LinearChainCRF
• 半马尔可夫 CRF
• 依存句法 CRF
• 非投影视图依存句法 CRF
• 树状 CRF
• 神经 PCFG / 神经 HMM

每个分布都包含：

• 最大值、采样、熵、归一化常数、掩码、对数概率、k 最大值

扩展功能：

• 与 torchtext、pytorch-transformers、dgl 的集成
• 生成式结构化模型（CFG / HMM / HSMM）的适配器
• 常见的树结构参数化方法，如 TreeLSTM / SpanLSTM

低层 API：

所有功能均通过半环动态规划实现。

• 对数边缘概率
• 最大值和 MAP 计算
• 通过专用反向传播进行采样
• 熵及一阶半环计算。

示例

• BERT 词性标注
• BERT 依存句法分析
• 无监督学习
• 结构化变分自编码器

引用

@misc{alex2020torchstruct,
    title={Torch-Struct: 深度结构化预测库},
    author={Alexander M. Rush},
    year={2020},
    eprint={2002.00876},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}

本研究部分得到了 NSF IIS-1901030 资助。

pytorch-struct 快速上手指南

pytorch-struct 是一个专为深度学习应用设计的结构化预测库，提供了经过测试且支持 GPU 加速的核心算法实现（如 HMM、CRF、依存句法树、PCFG 等）。它旨在作为高效的批处理层嵌入到其他 PyTorch 代码中。

环境准备

• 系统要求：Linux / macOS / Windows
• 前置依赖：
  • Python >= 3.6
  • PyTorch >= 1.4 (推荐最新稳定版)
  • CUDA (可选，用于 GPU 加速)
• 可选依赖：
  • genbmm：用于加速 FastLogSemiring 的 CUDA 内核。
  • matplotlib：用于可视化结果。

安装步骤

推荐使用 pip 直接从 GitHub 安装最新版本。国内用户若遇到连接问题，可配置 pip 使用国内镜像源（如清华源或阿里源）。

# 基础安装
pip install -qU git+https://github.com/harvardnlp/pytorch-struct

# (可选) 安装高性能 CUDA 内核 (需要编译环境)
pip install -qU git+https://github.com/harvardnlp/genbmm

# (可选) 安装绘图工具
pip install -q matplotlib

提示：若下载速度慢，可使用国内镜像加速： pip install -qU git+https://github.com/harvardnlp/pytorch-struct -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下示例演示如何构建一个依存句法 CRF (DependencyCRF) 模型，并计算边缘概率、最优解及采样。

1. 导入库与初始化

import torch
from torch_struct import DependencyCRF, LinearChainCRF
import matplotlib.pyplot as plt

def show(x): 
    """辅助函数：显示张量图像"""
    plt.imshow(x.detach().cpu())
    plt.show()

# 构造模拟数据 (Batch=2, Length=10, Length=10)
vals = torch.zeros(2, 10, 10) + 1e-5
vals[:, :5, :5] = torch.rand(5)
vals[:, 5:, 5:] = torch.rand(5) 

# 初始化分布 (输入通常为 log potentials)
dist = DependencyCRF(vals.log())

2. 核心功能演示

查看对数势函数：

show(dist.log_potentials[0])

计算边缘概率 (Marginals)：

show(dist.marginals[0])

计算最优解 (Argmax / Viterbi)：

show(dist.argmax.detach()[0])

计算配分函数与得分：

# 计算配分函数 (Forward / Inside algorithm)
log_partition = dist.partition

# 计算最优路径的对数概率
max_score = dist.log_prob(dist.argmax)

采样 (Sampling)：

# 采样一个结构
show(dist.sample((1,)).detach()[0, 0])

3. 处理变长序列 (Masking)

库内置了对 Padding 和 Masking 的支持，只需传入 lengths 参数。

# 指定序列长度 [10, 7]，第二个序列后 3 个位置将被自动掩码
dist = DependencyCRF(vals, lengths=torch.tensor([10, 7]))

# 观察不同长度的边缘概率
show(dist.marginals[0])
show(dist.marginals[1])

4. 其他结构支持

除了依存树，库还支持线性链 CRF (LinearChainCRF) 等多种结构：

# 构建线性链数据
chain = torch.zeros(2, 10, 10, 10) + 1e-5
chain[:, :, :, :] = vals.unsqueeze(-1).exp()
chain[:, :, :, :] += torch.eye(10, 10).view(1, 1, 10, 10) 
chain[:, 0, :, 0] = 1
chain[:, -1, 9, :] = 1
chain = chain.log()

# 使用 LinearChainCRF
dist_chain = LinearChainCRF(chain)
show(dist_chain.marginals.detach()[0].sum(-1))

更多高级用法（如结合 BERT、无监督学习、结构化 VAE 等）请参考官方 Notebook 示例。