【MIIT program】add miad model#291
Conversation
MiAD 精度对齐验证补充本文档为 pr.md 第 3 章(精度对齐验证)的补充,包含测试代码、数据和分析方法。 1. 前向精度对齐测试代码import paddle, numpy as np
from ppmat.models.miad.miad import MiAD
model_cfg = {
'hidden_dim': 512, 'latent_dim': 256, 'num_layers': 6,
'max_atoms': 100, 'act_fn': 'silu', 'dis_emb': 'sin',
'num_freqs': 128, 'edge_style': 'fc', 'cutoff': 6.0,
'max_neighbors': 20, 'ln': True, 'ip': True,
'smooth': False, 'pred_type': True,
}
diffusion_cfg = {
'method': 'DiffCSP', 'task': 'gen_mp20', 'cont_time': False,
'num_steps': 1000,
'lat_diffusion': {'method': 'ddpm', 'scheduler': 'diffcsp_cosine', 'cond_coef': 1.},
'frac_diffusion': {'method': 'wrapped_normal', 'scheduler': 'default_wrapped_normal'},
'type_diffusion': {'method': 'd3pm', 'scheduler': 'default_d3pm'},
}
def test_forward_alignment(batch2_atoms20_18, batch1_atoms12, batch3_atoms10_15_8):
"""
分别在 PyTorch 和 Paddle 上加载相同权重,输入相同 batch,
对比 lattice/coord/type 三个输出头的 mean absolute diff。
"""
model_pd = MiAD(model_cfg=model_cfg, diffusion_cfg=diffusion_cfg)
state_dict = paddle.load("pretrained-pd/miad_mp20_epoch8000.pdparams")
model_pd.set_state_dict(state_dict)
model_pd.eval()
results = []
for name, batch in [("batch2_atoms20_18", batch2_atoms20_18),
("batch1_atoms12", batch1_atoms12),
("batch3_atoms10_15_8", batch3_atoms10_15_8)]:
t = paddle.to_tensor([0.5])
out_pd = model_pd(batch) # Paddle forward
# PyTorch forward (same weights, same input)
out_pt = model_pt(batch) # 原始 PyTorch 模型
lattice_diff = paddle.mean(paddle.abs(out_pd[0] - out_pt[0]))
coord_diff = paddle.mean(paddle.abs(out_pd[1] - out_pt[1]))
type_diff = paddle.mean(paddle.abs(out_pd[2] - out_pt[2]))
results.append((name, lattice_diff, coord_diff, type_diff))
return results测试结果
2. 生成指标对齐自复现性测试import paddle, numpy as np, json, math
from ppmat.models.miad.miad import MiAD
from ppmat.models.miad.collate import create_sampling_batch
from ppmat.models.miad.type_diffusion import D3PM
# ---------- 指标计算 ----------
def compute_volume(lattice_matrix):
a, b, c = lattice_matrix
return abs(np.dot(a, np.cross(b, c)))
def iqr_filter(values, multiplier):
arr = np.array(values)
q1, q3 = np.percentile(arr, [25, 75])
iq = q3 - q1
mask = (arr >= q1 - multiplier * iq) & (arr <= q3 + multiplier * iq)
return arr[mask]
def compute_cv(values):
arr = np.array(values, dtype=float)
return float('inf') if len(arr) < 2 or arr.mean() == 0 else float(arr.std(ddof=1) / arr.mean())
# ---------- 主测试 ----------
NUM_SEEDS = 7 # 种子数
SAMPLES_PER_SEED = 500 # 每种子样本数
VOLUME_RANGE = (100.0, 1000.0) # 有效体积范围
IQR_MULTIPLIER = 3.0 # IQR 乘数
THRESHOLD_CV = 0.05 # 5% 阈值
def run_reproducibility_test():
model = MiAD(model_cfg=model_cfg, diffusion_cfg=diffusion_cfg)
state_dict = paddle.load("pretrained-pd/miad_mp20_epoch8000.pdparams")
model.set_state_dict(state_dict)
model.eval(); model.cuda()
# 生成数据
all_data = []
for seed_idx in range(NUM_SEEDS):
seed = 42 + seed_idx
paddle.seed(seed); np.random.seed(seed)
for batch_idx in range((SAMPLES_PER_SEED + 15) // 16):
bs = min(16, SAMPLES_PER_SEED - batch_idx * 16)
batch = create_sampling_batch(batch_size=bs, num_atoms=None, device='gpu')
with paddle.no_grad():
output = model.sample(batch_data=batch, num_inference_steps=1000)
for c in output['result']:
lat = c['lattice'].numpy() if hasattr(c['lattice'], 'numpy') else np.array(c['lattice'])
all_data.append({'volume': compute_volume(lat), 'num_atoms': c['num_atoms'], 'seed': seed})
# 跨 seed CV 计算
seed_data = {}
for d in all_data:
seed_data.setdefault(d['seed'], []).append(d)
metric_keys = ['filtered_atoms_mean', 'atoms_std_all', 'volume_mean_iqr',
'volume_median_iqr', 'volume_std_iqr', 'valid_volume_rate']
per_seed = {k: [] for k in metric_keys}
for sid in sorted(seed_data):
sd = seed_data[sid]
volumes = np.array([d['volume'] for d in sd])
atoms = np.array([d['num_atoms'] for d in sd])
valid_mask = (volumes >= VOLUME_RANGE[0]) & (volumes <= VOLUME_RANGE[1])
vv, va = volumes[valid_mask], atoms[valid_mask]
valid_rate = len(vv) / len(volumes)
if len(vv) == 0:
continue
iqr_v = iqr_filter(vv, IQR_MULTIPLIER)
iqr_a = va[np.isin(vv, iqr_v)]
if len(iqr_v) == 0:
continue
per_seed['filtered_atoms_mean'].append(float(iqr_a.mean()))
per_seed['atoms_std_all'].append(float(atoms.std(ddof=1)))
per_seed['volume_mean_iqr'].append(float(iqr_v.mean()))
per_seed['volume_median_iqr'].append(float(np.median(iqr_v)))
per_seed['volume_std_iqr'].append(float(iqr_v.std(ddof=1)))
per_seed['valid_volume_rate'].append(float(valid_rate))
print(f"{'Metric':<25} {'CV':>10} {'Result':>8}")
all_pass = True
for mk in metric_keys:
cv = compute_cv(per_seed[mk])
status = 'PASS' if cv < THRESHOLD_CV else 'FAIL'
if cv >= THRESHOLD_CV:
all_pass = False
print(f"{mk:<25} {cv*100:>9.2f}% {status:>8}")
print(f"\nOverall: {'ALL PASS' if all_pass else 'SOME FAILED'}")测试结果配置: 7 seeds x 500 samples, IQR=3.0, vrange=[100,1000]
3. 训练 loss 对齐测试设置
测试代码import torch, paddle, numpy as np
from ppmat.models.miad.miad import MiAD
# 训练循环 (PT/Paddle 各自跑 5 epochs)
def train_loop(model, framework='paddle'):
losses = []
for epoch in range(5):
epoch_losses = []
for batch_idx in range(10):
batch = sample_batch(epoch * 10 + batch_idx)
loss = model.train_step(batch)
loss.backward()
optimizer.step(); optimizer.zero_grad()
epoch_losses.append(float(loss))
losses.append(epoch_losses)
return losses
# PT 训练
pt_model, pt_opt = build_pt_model()
pt_losses = train_loop(pt_model, 'torch')
# Paddle 训练
pd_model = MiAD(model_cfg=cfg, diffusion_cfg=diffusion_cfg)
pd_model.set_state_dict(paddle.load("pretrained-pd/miad_mp20_epoch8000.pdparams"))
pd_opt = paddle.optimizer.Adam(parameters=pd_model.parameters(), learning_rate=1e-3)
pd_losses = train_loop(pd_model, 'paddle')
# 跨框架对比
for ep in range(5):
pt_avg = np.mean(pt_losses[ep])
pd_avg = np.mean(pd_losses[ep])
print(f"Epoch {ep}: PT avg={pt_avg:.3f}, PD avg={pd_avg:.3f}, |diff|={abs(pt_avg-pd_avg):.3f}")Epoch 级 loss 对比
Batch 级loss对比 (来自最后一个Epoch4)
|
|
请使用套件内已有的模块,并且参考指南 #258 |
1fa058e to
0a4e755
Compare
已经优化 |
leeleolay
left a comment
There was a problem hiding this comment.
miad是基于diffcsp的话,是否可以基于已有的一些实现,通过调用相关模块
| from ppmat.datasets.custom_data_type import ConcatNumpyWarper | ||
| from ppmat.datasets.geometric_data_type.batch import Batch | ||
| from ppmat.datasets.geometric_data_type.data import Data | ||
| from ppmat.models.miad.collate import MiADCollator # noqa: F401 |
| return energies | ||
|
|
||
|
|
||
| def relax_structure( |
| return _chgnet_model, _graph_converter | ||
|
|
||
|
|
||
| def predict_energy( |
| return None, None, 0 | ||
|
|
||
|
|
||
| def predict_energies_with_relaxation( |
| return str(sorted(list(structure.atomic_numbers))) | ||
|
|
||
|
|
||
| def _structures_from_cif_strings(cif_strings: List[str]) -> List[Optional[Structure]]: |
There was a problem hiding this comment.
是否可以直接调用build_structure
There was a problem hiding this comment.
参考ppmat/dataset/collate
There was a problem hiding this comment.
miad是基于diffcsp吗,是否可以直接使用相关实现
dc09ad3 to
2e914b7
Compare
按照各项说明,重新优化后提交 |
There was a problem hiding this comment.
修改为sun_metric_utils.py吧
There was a problem hiding this comment.
参考已有的ppmat/schdulers里的实现
| # MiAD | ||
|
|
||
| MiAD (Mirage Atom Diffusion) is a diffusion-based framework for de novo crystal generation. It introduces the concept of Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials. | ||
|
|
|
|
||
| MiAD (Mirage Atom Diffusion) is a diffusion-based framework for de novo crystal generation. It introduces the concept of Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials. | ||
|
|
||
| ## How It Works |
|
|
||
| 4. Post-Processing: Remaining mirage atoms are stripped before exporting to final CIF files . | ||
|
|
||
| ## Architecture |
| - paddle_scatter >= 2.1.2 | ||
| - numpy, pymatgen, omegaconf | ||
|
|
||
| ## Checkpoints |
There was a problem hiding this comment.
删除,并提供链接和result,在readme里补充result
|
提供预训练模型权重 |
…-model # Conflicts: # ppmat/metrics/__init__.py # ppmat/models/__init__.py
1524abb to
b9b34b8
Compare
|
代码做了修改:
文档部分做了修改:
|
| ## Dataset | ||
|
|
||
| | Source | Link | | ||
| |-----------------------|------| | ||
| | Original Google Drive | [Google Drive](https://drive.google.com/file/d/1BLI3VtvzfIIXlH6UHQ4o-gQaCIOZ1UR7/view?usp=sharing) | | ||
| | Mirror AiStudio | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) | | ||
|
|
||
| Extract to `./data/mp_20/` so that CSV files are at `./data/mp_20/train.csv`, `./data/mp_20/val.csv`, and `./data/mp_20/test.csv`. |
There was a problem hiding this comment.
| ## Results | ||
|
|
||
| | Metric | CHGNet | eq-V2 | | ||
| |--------|--------|-------| | ||
| | S.U.N. | 12.21% | 5.64% | |
There was a problem hiding this comment.
不是该模型的results,参考其他readme的内容
| ### Checkpoints | ||
|
|
||
| | Source | Link | | ||
| |--------------------------------|------| | ||
| | Original (PyTorch format) | [Google Drive](https://drive.google.com/file/d/1KyD6KzvjYFPfU8lutFyO_0b8EbeHSGqf/view?usp=sharing) | | ||
| | Mirror (PyTorch format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) | | ||
| | PaddleMaterials(Paddle format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48638/intro) | |
|
|
||
| MiAD introduces Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory. By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials. It uses DiffCsp as the backbone denoising architecture. | ||
|
|
||
|  |
There was a problem hiding this comment.
下载该文件放到该任务下的docs里,并修改链接
|
|
||
| --- | ||
|
|
||
| ## Model Description |
There was a problem hiding this comment.
这里没有配置metric,需兼容已有的metric实现方式
There was a problem hiding this comment.
兼容已有metric的实现方式,建议绑定streaming
|
structure_generation/README.md 在该页面补充模型信息 |
|
放弃了之前的 1:1的代码转译的paconvert的代码模式; 重新组织和复用了代码。
这个markdown 我感觉需要后续集中一次修改,不要每个提交都修改,我尝试修改了一下;发现这个表格是 每次增加一列的。就意味着 前后2个分支提交,一定会冲突。 |
MiAD 模型迁移
1. 概述
将 MiAD (Mirage Atom Diffusion) 晶体生成模型从 PyTorch 框架迁移至 PaddlePaddle 框架。
2. 权重信息
2.1 原版权重
miad_mp20_epoch8000.pt原版权重不方便下载的,已经镜像到aistudio上, https://aistudio.baidu.com/modelsdetail/48578
2.2 转换权重
权重已成功转换为 PaddlePaddle 格式,存储在
miad_mp20_epoch8000.pdparams:已经发布到aistudio上, https://aistudio.baidu.com/modelsdetail/48638
或者通过http直接下载 miad_mp20.zip
paddle代码已经内置了默认下载地址,使用
structure_generation/sample.py可直接加载miad_mp20中的yaml配置和预训练权重。3. 精度对齐验证
3.1 前向精度对齐
测试设置
测试结果
3.2 生成指标对齐
测试设置
多次采样误差
3.3 训练 loss 对齐
测试设置
Epoch 级loss对比
Batch 级loss对比 (来自最后一个Epoch4)
两边的数据加载器和扩散采样使用各自独立的 RNG,导致输入数据本身不同,最终loss有一定差距
为了避免GitHub的PR提交64K最大限制,剩余脚本在下面的回复中。