OpenAlphaDiffract — 昇腾 NPU 适配版

OpenAlphaDiffract 是面向粉末 X 射线衍射（PXRD）数据的自动晶体学分析模型。本仓库为 昇腾 NPU（Ascend910）适配版本，完成了从 CUDA/PyTorch 到 torch_npu 的迁移，并附带一套自包含的推理服务。模型在 NPU 与 CPU 上输出完全一致，单条 8192 点图谱推理时延约 1.4 ms，单卡即可流畅承载在线分析请求。

简介

模型同时完成三个任务，输入单条 PXRD 图谱，一次前向即可得到全部结果：

• 架构：1D ConvNeXt 主干（3 级下采样，560 维扁平特征）+ 三个 MLP 头
• 参数量：8,734,989（约 8.7M）
• 权重：model.safetensors（35 MB，41 个张量，float32）
• 输入约定：8192 点强度向量，负值清零后线性归一化到 [0, 100]
• 原仓：AdvancedPhotonSource/OpenAlphaDiffract，论文见 arXiv:2603.23367

验证环境

本次适配与验证使用的具体环境如下：

服务启动

方式一：直接启动推理服务（推荐）

pip install -r requirements.txt
python inference.py --serve --model /path/to/model_dir --host 0.0.0.0 --port 8100

其中 /path/to/model_dir 需包含 config.json、model.safetensors、maxsub.json。服务默认自动选择 npu:0 设备，启动后监听在指定端口。

方式二：命令行单次推理

python inference.py --model /path/to/model_dir --pattern /path/to/pattern.npy

--device 可显式指定 npu:0 / cuda:0 / cpu，缺省自动探测。

Smoke 验证

服务启动后，按以下顺序做最小验证：

# 1. 健康检查
curl -s http://127.0.0.1:8100/health
# 期望: {"status":"ok","device":"npu:0","model":"AlphaDiffract"}

# 2. 模型列表（OpenAI 兼容）
curl -s http://127.0.0.1:8100/v1/models

# 3. 单条预测
curl -s -X POST http://127.0.0.1:8100/v1/predict \
  -H 'Content-Type: application/json' \
  -d '{"pattern":[0.1,0.2,...], "top_k_sg":3}'

实际验证输出（示例图谱，top_k_sg=3）：

{
  "crystal_system": {"index": 0, "name": "Triclinic", "probability": 0.7073},
  "space_group": {"best": {"space_group": 1, "probability": 0.3099}},
  "lattice_parameters": [
    {"parameter": "a", "value": 6.9728, "unit": "A"},
    {"parameter": "b", "value": 8.6647, "unit": "A"},
    {"parameter": "c", "value": 11.5285, "unit": "A"},
    {"parameter": "alpha", "value": 95.3637, "unit": "deg"},
    {"parameter": "beta", "value": 93.7262, "unit": "deg"},
    {"parameter": "gamma", "value": 95.6381, "unit": "deg"}
  ]
}

性能参考

以下数据均在上述 Ascend910 单卡环境实测，batch=1 时延为 100 次平均。

时延与吞吐数据说明：模型本身为轻量 1D CNN，计算量小，瓶颈在 H2D 拷贝与框架调度开销；吞吐在 batch=32 后趋于饱和。

精度测评

采用「CPU 基线 + NPU 结果」逐样本比对的方式评估迁移精度。对 30 条多样化合成图谱（含干净/噪声/锐峰三类），结果如下：

分类任务在 NPU 与 CPU 上完全一致；回归输出的微小差异属于 float32 计算舍入，量级远小于模型本身的预测误差，不影响可用性。

注意事项

• 本版本针对 单模型（Materials Project 数据训练）适配，非论文中的 10 模型集成，性能参考请以论文为准。
• maxsub.json 为 GEMD 空间群子群图（230×230 LUT），推理时必须可用；若缺失，加载逻辑会尝试从权重恢复 gemd_distance_matrix buffer。
• 输入长度固定为 8192；脚本会自动截断/补零对齐，但强烈建议上游归一化到 [0, 100]。
• 服务以异步 worker 处理请求，模型推理本身在 no_grad 下进行，权重为 float32，未做量化。
• 更换 NPU 型号或 CANN 版本时，请重新跑一次 Smoke 验证 中的三个请求确认输出结构一致。
• 若 8100 端口被占用，请改用其它空闲端口，勿抢占其他服务。