已开启
"add boltzgen NPU adaptation" #19
Hugh创建于 8月7日
"add boltzgen NPU adaptation" #19
已开启
共 4 个文件变更+524-0
| @@ -0,0 +1,82 @@ | |||
| 1 | +--- | ||
| 2 | +license: mit | ||
| 3 | +--- | ||
| 4 | + | ||
| 5 | +# 概述 | ||
| 6 | + | ||
| 7 | +BoltzGen 是一种基于扩散模型的生成式蛋白质设计与逆折叠方法。它通过扩散生成过程产生蛋白质序列候选，结合逆折叠模型和折叠置信度评分进行排序筛选，能够在给定骨架结构条件下生成多样且高质量的蛋白质序列设计。该方法集成了设计（design）、逆折叠（inverse folding）、折叠（folding）、分析（analysis）和过滤（filtering）完整 pipeline。 | ||
| 8 | + | ||
| 9 | +- 参考实现： | ||
| 10 | + | ||
| 11 | + ```shell | ||
| 12 | + url=https://github.com/HannesStark/boltzgen.git | ||
| 13 | + commit_id=a3149cf18eeb58648d1abbb27539bd73f746cdda | ||
| 14 | + ``` | ||
| 15 | + | ||
| 16 | +- 适配昇腾 AI 处理器的实现： | ||
| 17 | + | ||
| 18 | + ```shell | ||
| 19 | + url=https://gitcode.com/AI4Science/LifeScience.git | ||
| 20 | + code_path=PyTorch/boltzgen | ||
| 21 | + ``` | ||
| 22 | + | ||
| 23 | +#### 组件版本 | ||
| 24 | + | ||
| 25 | +```shell | ||
| 26 | +cann：8.5.1 | ||
| 27 | +python：3.11 | ||
| 28 | +torch：2.7.1 | ||
| 29 | +torch-npu：2.7.1 | ||
| 30 | +``` | ||
| 31 | + | ||
| 32 | +#### 拉取模型代码 | ||
| 33 | + | ||
| 34 | +```shell | ||
| 35 | +git clone https://gitcode.com/AI4Science/LifeScience.git | ||
| 36 | +cd LifeScience/PyTorch/boltzgen | ||
| 37 | +git clone https://github.com/HannesStark/boltzgen.git && cd boltzgen | ||
| 38 | +git checkout a3149cf18eeb58648d1abbb27539bd73f746cdda | ||
| 39 | +git apply ../patch/boltzgen.patch | ||
| 40 | +``` | ||
| 41 | + | ||
| 42 | +#### 环境准备 | ||
| 43 | + | ||
| 44 | +1、创建 conda 环境 | ||
| 45 | +``` | ||
| 46 | +conda create -n boltzgen_npu python=3.11 -y | ||
| 47 | +conda activate boltzgen_npu | ||
| 48 | +``` | ||
| 49 | +2、设置 CANN 环境变量 | ||
| 50 | +``` | ||
| 51 | +source /usr/local/Ascend/ascend-toolkit/set_env.sh | ||
| 52 | +``` | ||
| 53 | +3、安装 requirements.txt | ||
| 54 | +``` | ||
| 55 | +pip install -r boltzgen_requirements.txt | ||
| 56 | +``` | ||
| 57 | +4、安装 boltzgen | ||
| 58 | +``` | ||
| 59 | +cd /home/workspace/boltzgen && pip install -e . | ||
| 60 | +``` | ||
| 61 | + | ||
| 62 | +#### 模型权重准备 | ||
| 63 | + | ||
| 64 | +模型权重文件需预先放置到 `~/.cache/` 目录下（如机器可访问外网则会自动从 HuggingFace 下载）： | ||
| 65 | +- `boltzgen1_diverse.ckpt` — 扩散设计模型（多样性） | ||
| 66 | +- `boltzgen1_adherence.ckpt` — 扩散设计模型（贴合度） | ||
| 67 | +- `boltzgen1_ifold.ckpt` — 反向折叠模型 | ||
| 68 | +- `boltz2_conf_final.ckpt` — 折叠/置信度模型 | ||
| 69 | +- `mols.zip` — 分子库（CCD 数据） | ||
| 70 | + | ||
| 71 | +权重来源：HuggingFace `boltzgen/boltzgen-1` 和 `boltzgen/inference-data`。 | ||
| 72 | + | ||
| 73 | +#### 数据集准备 | ||
| 74 | +``` | ||
| 75 | +使用模型自带用例验证，位置是：boltzgen/example/vanilla_protein/1g13prot.yaml，若想使用其他数据自行准备 yaml 配置即可。 | ||
| 76 | +``` | ||
| 77 | + | ||
| 78 | +#### 推理 | ||
| 79 | +模型首次推理时间较长，请耐心等待。 | ||
| 80 | +``` | ||
| 81 | +bash test/inference.sh | ||
| 82 | +``` | ||
| @@ -0,0 +1,17 @@ | |||
| 1 | +# BoltzGen 昇腾 NPU requirements.txt | ||
| 2 | +# 环境：CANN Toolkit 8.3.RC1 / Linux aarch64 / Python 3.11 | ||
| 3 | +# 安装顺序： | ||
| 4 | +# 1) source /usr/local/Ascend/ascend-toolkit/set_env.sh | ||
| 5 | +# 2) pip install -r requirements.txt | ||
| 6 | +# 3) cd /home/workspace/boltzgen && pip install -e . | ||
| 7 | +torch==2.7.1 | ||
| 8 | +torch_npu==2.7.1 | ||
| 9 | +numpy==2.0.2 | ||
| 10 | +decorator==5.3.1 | ||
| 11 | +attrs==26.1.0 | ||
| 12 | +psutil==7.2.2 | ||
| 13 | +absl-py==2.5.0 | ||
| 14 | +cloudpickle==3.1.2 | ||
| 15 | +ml_dtypes==0.5.4 | ||
| 16 | +scipy==1.17.1 | ||
| 17 | +tornado==6.5.7 | ||
| @@ -0,0 +1,293 @@ | |||
| 1 | +diff --git a/pyproject.toml b/pyproject.toml | ||
| 2 | +index 6e32f4b..98464e4 100755 | ||
| 3 | +--- a/pyproject.toml | ||
| 4 | ++++ b/pyproject.toml | ||
| 5 | + dependencies = [ | ||
| 6 | + "einx", | ||
| 7 | + "einops", | ||
| 8 | + "mashumaro", | ||
| 9 | +- "nvidia-ml-py>=12.535.133", | ||
| 10 | +- "cuequivariance_ops_cu12>=0.5.0", | ||
| 11 | +- "cuequivariance_ops_torch_cu12>=0.5.0", | ||
| 12 | +- "cuequivariance_torch>=0.5.0", | ||
| 13 | ++# "nvidia-ml-py>=12.535.133", | ||
| 14 | ++# "cuequivariance_ops_cu12>=0.5.0", | ||
| 15 | ++# "cuequivariance_ops_torch_cu12>=0.5.0", | ||
| 16 | ++# "cuequivariance_torch>=0.5.0", | ||
| 17 | + "huggingface_hub", | ||
| 18 | + "biopython", | ||
| 19 | + ] | ||
| 20 | +diff --git a/src/boltzgen/cli/boltzgen.py b/src/boltzgen/cli/boltzgen.py | ||
| 21 | +index 4cf6c6c..0a4a874 100644 | ||
| 22 | +--- a/src/boltzgen/cli/boltzgen.py | ||
| 23 | ++++ b/src/boltzgen/cli/boltzgen.py | ||
| 24 | + | ||
| 25 | ++try: | ||
| 26 | ++ import torch_npu | ||
| 27 | ++ from torch_npu.contrib import transfer_to_npu | ||
| 28 | ++ import torch | ||
| 29 | ++ _orig_get_device_capability = torch.cuda.get_device_capability | ||
| 30 | ++ def _patched_get_device_capability(device=None): | ||
| 31 | ++ cap = _orig_get_device_capability(device) | ||
| 32 | ++ return cap if cap is not None else (9, 0) | ||
| 33 | ++ torch.cuda.get_device_capability = _patched_get_device_capability | ||
| 34 | ++except ImportError: | ||
| 35 | ++ pass | ||
| 36 | ++ | ||
| 37 | + #!/usr/bin/env python3 | ||
| 38 | + """ | ||
| 39 | + This script orchestrates work. It sets up an output directory with yaml files of pipeline steps that need to be run, and launches processes that run the pipeline steps. | ||
| 40 | + def get_artifact_path( | ||
| 41 | + raise ValueError( | ||
| 42 | + f"Invalid artifact: {artifact}. Expected format: huggingface:<repo_id>:<filename>" | ||
| 43 | + ) | ||
| 44 | +- result = huggingface_hub.hf_hub_download( | ||
| 45 | +- repo_id, | ||
| 46 | +- filename, | ||
| 47 | +- repo_type=repo_type, | ||
| 48 | +- library_name="boltzgen", | ||
| 49 | +- force_download=args.force_download, | ||
| 50 | +- token=args.models_token, | ||
| 51 | +- cache_dir=args.cache, | ||
| 52 | +- ) | ||
| 53 | +- result = Path(result) | ||
| 54 | ++ | ||
| 55 | ++ cache_dir = args.cache or Path.home() / ".cache" | ||
| 56 | ++ local_path = Path(cache_dir) / filename | ||
| 57 | ++ if local_path.exists() and not args.force_download: | ||
| 58 | ++ result = local_path | ||
| 59 | ++ else: | ||
| 60 | ++ result = huggingface_hub.hf_hub_download( | ||
| 61 | ++ repo_id, | ||
| 62 | ++ filename, | ||
| 63 | ++ repo_type=repo_type, | ||
| 64 | ++ library_name="boltzgen", | ||
| 65 | ++ force_download=args.force_download, | ||
| 66 | ++ token=args.models_token, | ||
| 67 | ++ cache_dir=args.cache, | ||
| 68 | ++ ) | ||
| 69 | ++ result = Path(result) | ||
| 70 | ++ | ||
| 71 | ++# result = huggingface_hub.hf_hub_download( | ||
| 72 | ++# repo_id, | ||
| 73 | ++# filename, | ||
| 74 | ++# repo_type=repo_type, | ||
| 75 | ++# library_name="boltzgen", | ||
| 76 | ++# force_download=args.force_download, | ||
| 77 | ++# token=args.models_token, | ||
| 78 | ++# cache_dir=args.cache, | ||
| 79 | ++# ) | ||
| 80 | ++# result = Path(result) | ||
| 81 | + else: | ||
| 82 | + result = Path(artifact) | ||
| 83 | + if not result.exists(): | ||
| 84 | +diff --git a/src/boltzgen/model/layers/npu_kernels.py b/src/boltzgen/model/layers/npu_kernels.py | ||
| 85 | +new file mode 100644 | ||
| 86 | +index 0000000..0c629f3 | ||
| 87 | +--- /dev/null | ||
| 88 | ++++ b/src/boltzgen/model/layers/npu_kernels.py | ||
| 89 | + | ||
| 90 | ++"""NPU-native implementations of cuequivariance_torch operations. | ||
| 91 | ++ | ||
| 92 | ++Provides functionally equivalent replacements for: | ||
| 93 | ++ 1. cuequivariance_torch.primitives.triangle.triangle_attention | ||
| 94 | ++ 2. cuequivariance_torch.primitives.triangle.triangle_multiplicative_update | ||
| 95 | ++ | ||
| 96 | ++These run entirely on Ascend NPU using standard PyTorch ops (no CPU fallback). | ||
| 97 | ++""" | ||
| 98 | ++ | ||
| 99 | ++import math | ||
| 100 | ++ | ||
| 101 | ++import torch | ||
| 102 | ++import torch.nn.functional as F | ||
| 103 | ++from torch import Tensor | ||
| 104 | ++ | ||
| 105 | ++ | ||
| 106 | ++def npu_triangle_attention( | ||
| 107 | ++ q: Tensor, | ||
| 108 | ++ k: Tensor, | ||
| 109 | ++ v: Tensor, | ||
| 110 | ++ tri_bias: Tensor, | ||
| 111 | ++ mask: Tensor, | ||
| 112 | ++ scale: float, | ||
| 113 | ++) -> Tensor: | ||
| 114 | ++ """NPU-native triangle attention. | ||
| 115 | ++ | ||
| 116 | ++ Equivalent to cuequivariance_torch.primitives.triangle.triangle_attention. | ||
| 117 | ++ | ||
| 118 | ++ Parameters | ||
| 119 | ++ ---------- | ||
| 120 | ++ q : Tensor [*, H, Q, C] | ||
| 121 | ++ k : Tensor [*, H, K, C] | ||
| 122 | ++ v : Tensor [*, H, K, C] | ||
| 123 | ++ tri_bias : Tensor [*, 1, H, Q, K] or [*, H, Q, K] | ||
| 124 | ++ mask : Tensor bool [*, Q, K] or broadcastable | ||
| 125 | ++ scale : float | ||
| 126 | ++ | ||
| 127 | ++ Returns | ||
| 128 | ++ ------- | ||
| 129 | ++ Tensor [*, H, Q, C] | ||
| 130 | ++ """ | ||
| 131 | ++ q = q * scale | ||
| 132 | ++ | ||
| 133 | ++ attn = torch.matmul(q, k.transpose(-1, -2)) | ||
| 134 | ++ | ||
| 135 | ++ if tri_bias.dim() > attn.dim(): | ||
| 136 | ++ tri_bias = tri_bias.squeeze(-4) | ||
| 137 | ++ attn = attn + tri_bias | ||
| 138 | ++ | ||
| 139 | ++ if mask is not None: | ||
| 140 | ++ if mask.dim() < attn.dim(): | ||
| 141 | ++ mask = mask.unsqueeze(-3) | ||
| 142 | ++ attn = attn.masked_fill(~mask, float("-inf")) | ||
| 143 | ++ | ||
| 144 | ++ attn = F.softmax(attn, dim=-1) | ||
| 145 | ++ out = torch.matmul(attn, v) | ||
| 146 | ++ return out | ||
| 147 | ++ | ||
| 148 | ++ | ||
| 149 | ++def npu_triangle_multiplicative_update( | ||
| 150 | ++ x: Tensor, | ||
| 151 | ++ direction: str, | ||
| 152 | ++ mask: Tensor, | ||
| 153 | ++ norm_in_weight: Tensor, | ||
| 154 | ++ norm_in_bias: Tensor, | ||
| 155 | ++ p_in_weight: Tensor, | ||
| 156 | ++ g_in_weight: Tensor, | ||
| 157 | ++ norm_out_weight: Tensor, | ||
| 158 | ++ norm_out_bias: Tensor, | ||
| 159 | ++ p_out_weight: Tensor, | ||
| 160 | ++ g_out_weight: Tensor, | ||
| 161 | ++ eps: float = 1e-5, | ||
| 162 | ++) -> Tensor: | ||
| 163 | ++ """NPU-native triangle multiplicative update. | ||
| 164 | ++ | ||
| 165 | ++ Equivalent to cuequivariance_torch.primitives.triangle.triangle_multiplicative_update. | ||
| 166 | ++ | ||
| 167 | ++ Parameters | ||
| 168 | ++ ---------- | ||
| 169 | ++ x : Tensor [B, N, N, D] | ||
| 170 | ++ direction : str "outgoing" or "incoming" | ||
| 171 | ++ mask : Tensor [B, N, N] | ||
| 172 | ++ norm_in_weight, norm_in_bias : LayerNorm parameters | ||
| 173 | ++ p_in_weight : Tensor [2D, D] projection weight | ||
| 174 | ++ g_in_weight : Tensor [2D, D] gating weight | ||
| 175 | ++ norm_out_weight, norm_out_bias : LayerNorm parameters | ||
| 176 | ++ p_out_weight : Tensor [D, D] output projection weight | ||
| 177 | ++ g_out_weight : Tensor [D, D] output gating weight | ||
| 178 | ++ eps : float | ||
| 179 | ++ | ||
| 180 | ++ Returns | ||
| 181 | ++ ------- | ||
| 182 | ++ Tensor [B, N, N, D] | ||
| 183 | ++ """ | ||
| 184 | ++ x_normed = F.layer_norm(x, [x.shape[-1]], norm_in_weight, norm_in_bias, eps) | ||
| 185 | ++ x_in = x_normed | ||
| 186 | ++ | ||
| 187 | ++ proj = F.linear(x_normed, p_in_weight) | ||
| 188 | ++ gate = F.linear(x_normed, g_in_weight).sigmoid() | ||
| 189 | ++ x_gated = proj * gate | ||
| 190 | ++ | ||
| 191 | ++ x_gated = x_gated * mask.unsqueeze(-1) | ||
| 192 | ++ | ||
| 193 | ++ a, b = x_gated.float().chunk(2, dim=-1) | ||
| 194 | ++ | ||
| 195 | ++ if direction == "outgoing": | ||
| 196 | ++ x_tri = torch.einsum("bikd,bjkd->bijd", a, b) | ||
| 197 | ++ elif direction == "incoming": | ||
| 198 | ++ x_tri = torch.einsum("bkid,bkjd->bijd", a, b) | ||
| 199 | ++ else: | ||
| 200 | ++ raise ValueError(f"direction must be outgoing or incoming, got {direction}") | ||
| 201 | ++ | ||
| 202 | ++ x_out = F.layer_norm(x_tri, [x_tri.shape[-1]], norm_out_weight, norm_out_bias, eps) | ||
| 203 | ++ x_out = F.linear(x_out, p_out_weight) | ||
| 204 | ++ | ||
| 205 | ++ gate_out = F.linear(x_in, g_out_weight).sigmoid() | ||
| 206 | ++ return x_out * gate_out | ||
| 207 | + | ||
| 208 | +diff --git a/src/boltzgen/model/layers/triangular.py b/src/boltzgen/model/layers/triangular.py | ||
| 209 | +index e7b0817..c43c238 100755 | ||
| 210 | +--- a/src/boltzgen/model/layers/triangular.py | ||
| 211 | ++++ b/src/boltzgen/model/layers/triangular.py | ||
| 212 | + def _kernel_triangular_mult( | ||
| 213 | + g_out_weight: Tensor, | ||
| 214 | + eps: float, | ||
| 215 | + ): | ||
| 216 | ++ | ||
| 217 | ++ if x.device.type == "npu": | ||
| 218 | ++ from boltzgen.model.layers.npu_kernels import npu_triangle_multiplicative_update | ||
| 219 | ++ return npu_triangle_multiplicative_update( | ||
| 220 | ++ x, direction=direction, mask=mask, | ||
| 221 | ++ norm_in_weight=norm_in_weight, norm_in_bias=norm_in_bias, | ||
| 222 | ++ p_in_weight=p_in_weight, g_in_weight=g_in_weight, | ||
| 223 | ++ norm_out_weight=norm_out_weight, norm_out_bias=norm_out_bias, | ||
| 224 | ++ p_out_weight=p_out_weight, g_out_weight=g_out_weight, eps=eps, | ||
| 225 | ++ ) | ||
| 226 | ++ | ||
| 227 | + try: | ||
| 228 | + from cuequivariance_torch.primitives.triangle import ( | ||
| 229 | + triangle_multiplicative_update as _triangle_multiplicative_update, | ||
| 230 | +diff --git a/src/boltzgen/model/layers/triangular_attention/primitives.py b/src/boltzgen/model/layers/triangular_attention/primitives.py | ||
| 231 | +index 60ca1f1..806fe69 100755 | ||
| 232 | +--- a/src/boltzgen/model/layers/triangular_attention/primitives.py | ||
| 233 | ++++ b/src/boltzgen/model/layers/triangular_attention/primitives.py | ||
| 234 | + def _attention( | ||
| 235 | + | ||
| 236 | + @torch.compiler.disable | ||
| 237 | + def kernel_triangular_attn(q, k, v, tri_bias, mask, scale): | ||
| 238 | ++ if q.device.type == 'npu': | ||
| 239 | ++ from boltzgen.model.layers.npu_kernels import npu_triangle_attention | ||
| 240 | ++ return npu_triangle_attention(q,k,v,tri_bias,mask=mask, scale=scale) | ||
| 241 | + from cuequivariance_torch.primitives.triangle import triangle_attention | ||
| 242 | + return triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale) | ||
| 243 | + | ||
| 244 | +diff --git a/src/boltzgen/resources/main.py b/src/boltzgen/resources/main.py | ||
| 245 | +index c0c486b..770bc21 100755 | ||
| 246 | +--- a/src/boltzgen/resources/main.py | ||
| 247 | ++++ b/src/boltzgen/resources/main.py | ||
| 248 | + | ||
| 249 | ++try: | ||
| 250 | ++ import torch_npu | ||
| 251 | ++ from torch_npu.contrib import transfer_to_npu | ||
| 252 | ++ import torch | ||
| 253 | ++ _orig_get_device_capability = torch.cuda.get_device_capability | ||
| 254 | ++ def _patched_get_device_capability(device=None): | ||
| 255 | ++ cap = _orig_get_device_capability(device) | ||
| 256 | ++ return cap if cap is not None else (9, 0) | ||
| 257 | ++ torch.cuda.get_device_capability = _patched_get_device_capability | ||
| 258 | ++except ImportError: | ||
| 259 | ++ pass | ||
| 260 | ++ | ||
| 261 | + import sys | ||
| 262 | + from typing import List | ||
| 263 | + import hydra | ||
| 264 | +diff --git a/src/boltzgen/task/analyze/analyze.py b/src/boltzgen/task/analyze/analyze.py | ||
| 265 | +index 256696a..d78e11a 100755 | ||
| 266 | +--- a/src/boltzgen/task/analyze/analyze.py | ||
| 267 | ++++ b/src/boltzgen/task/analyze/analyze.py | ||
| 268 | + class Analyze(Task): | ||
| 269 | + ).bool() | ||
| 270 | + # Use the same features as refolded, just need to change the coordinates back to the original designed coordinates | ||
| 271 | + feat_design["coords"] = feat_design["input_coords"] | ||
| 272 | ++ _folded_res = torch.from_numpy(folded["res_type"].squeeze()) | ||
| 273 | + if ( | ||
| 274 | +- not len(folded["res_type"].squeeze()) | ||
| 275 | +- == len(feat_design["res_type"]) | ||
| 276 | +- or not ( | ||
| 277 | +- folded["res_type"].squeeze() == feat_design["res_type"] | ||
| 278 | +- ).all() | ||
| 279 | ++ _folded_res.shape != feat_design["res_type"].shape | ||
| 280 | ++ or not (_folded_res == feat_design["res_type"]).all() | ||
| 281 | + ): | ||
| 282 | ++ | ||
| 283 | ++# if ( | ||
| 284 | ++# not len(folded["res_type"].squeeze()) | ||
| 285 | ++# == len(feat_design["res_type"]) | ||
| 286 | ++# or not ( | ||
| 287 | ++# folded["res_type"].squeeze() == feat_design["res_type"] | ||
| 288 | ++# ).all() | ||
| 289 | ++# ): | ||
| 290 | ++ | ||
| 291 | + msg = f"Skipping {path}. The sequences for which the refolding was run are not the same as the sequences in the design_dir. Maybe the designs in the design_dir were overwritten. Or maybe two processes are operating on the same design_dir." | ||
| 292 | + print(msg) | ||
| 293 | + return None | ||
| @@ -0,0 +1,132 @@ | |||
| 1 | +#!/bin/bash | ||
| 2 | + | ||
| 3 | +# ============================================ | ||
| 4 | +# Boltzgen NPU 推理脚本 | ||
| 5 | +# 功能：在昇腾NPU上运行蛋白质设计推理 | ||
| 6 | +# 作者：Auto Generated | ||
| 7 | +# 日期：2026-08-06 | ||
| 8 | +# ============================================ | ||
| 9 | + | ||
| 10 | +# 设置脚本在出错时退出 | ||
| 11 | +set -e | ||
| 12 | + | ||
| 13 | +# 定义颜色输出（便于查看日志） | ||
| 14 | +RED='\033[0;31m' | ||
| 15 | +GREEN='\033[0;32m' | ||
| 16 | +YELLOW='\033[1;33m' | ||
| 17 | +NC='\033[0m' # No Color | ||
| 18 | + | ||
| 19 | +# 打印带颜色的信息 | ||
| 20 | +print_info() { | ||
| 21 | + echo -e "${GREEN}[INFO]${NC} $1" | ||
| 22 | +} | ||
| 23 | + | ||
| 24 | +print_error() { | ||
| 25 | + echo -e "${RED}[ERROR]${NC} $1" | ||
| 26 | +} | ||
| 27 | + | ||
| 28 | +print_warning() { | ||
| 29 | + echo -e "${YELLOW}[WARNING]${NC} $1" | ||
| 30 | +} | ||
| 31 | + | ||
| 32 | +# ============================================ | ||
| 33 | +# 1. 环境配置 | ||
| 34 | +# ============================================ | ||
| 35 | + | ||
| 36 | +print_info "开始配置环境..." | ||
| 37 | + | ||
| 38 | +# 激活 conda 环境 | ||
| 39 | +print_info "激活 conda 环境: boltzgen_npu" | ||
| 40 | +conda activate boltzgen_npu || { | ||
| 41 | + print_error "conda 环境激活失败，请检查环境名称" | ||
| 42 | + exit 1 | ||
| 43 | +} | ||
| 44 | + | ||
| 45 | +# 设置 Ascend 环境变量 | ||
| 46 | +print_info "设置 Ascend 工具包环境变量" | ||
| 47 | +if [ -f "/usr/local/Ascend/ascend-toolkit/set_env.sh" ]; then | ||
| 48 | + source /usr/local/Ascend/ascend-toolkit/set_env.sh | ||
| 49 | + print_info "Ascend 环境变量设置成功" | ||
| 50 | +else | ||
| 51 | + print_error "Ascend 环境变量文件不存在: /usr/local/Ascend/ascend-toolkit/set_env.sh" | ||
| 52 | + exit 1 | ||
| 53 | +fi | ||
| 54 | + | ||
| 55 | +# 切换到工作目录 | ||
| 56 | +WORK_DIR="/home/workspace/boltzgen" | ||
| 57 | +print_info "切换到工作目录: $WORK_DIR" | ||
| 58 | +cd $WORK_DIR || { | ||
| 59 | + print_error "工作目录不存在: $WORK_DIR" | ||
| 60 | + exit 1 | ||
| 61 | +} | ||
| 62 | + | ||
| 63 | +# ============================================ | ||
| 64 | +# 2. 运行参数配置 | ||
| 65 | +# ============================================ | ||
| 66 | + | ||
| 67 | +# 可以在这里修改参数，而不必修改下面的命令 | ||
| 68 | +YAML_FILE="example/vanilla_protein/1g13prot.yaml" | ||
| 69 | +OUTPUT_DIR="workbench/test_run" | ||
| 70 | +PROTOCOL="protein-anything" | ||
| 71 | +NUM_DESIGNS=100 | ||
| 72 | +BUDGET=5 | ||
| 73 | +DEVICES=1 | ||
| 74 | +NUM_WORKERS=0 | ||
| 75 | + | ||
| 76 | +print_info "运行参数：" | ||
| 77 | +print_info " - YAML文件: $YAML_FILE" | ||
| 78 | +print_info " - 输出目录: $OUTPUT_DIR" | ||
| 79 | +print_info " - 协议: $PROTOCOL" | ||
| 80 | +print_info " - 设计数量: $NUM_DESIGNS" | ||
| 81 | +print_info " - 预算: $BUDGET" | ||
| 82 | +print_info " - 设备数: $DEVICES" | ||
| 83 | +print_info " - 工作进程: $NUM_WORKERS" | ||
| 84 | + | ||
| 85 | +# ============================================ | ||
| 86 | +# 3. 执行推理 | ||
| 87 | +# ============================================ | ||
| 88 | + | ||
| 89 | +print_info "开始运行 Boltzgen 推理..." | ||
| 90 | +print_info "开始时间: $(date '+%Y-%m-%d %H:%M:%S')" | ||
| 91 | + | ||
| 92 | +# 记录开始时间 | ||
| 93 | +START_TIME=$(date +%s) | ||
| 94 | + | ||
| 95 | +# 执行主命令 | ||
| 96 | +boltzgen run "$YAML_FILE" \ | ||
| 97 | + --output "$OUTPUT_DIR" \ | ||
| 98 | + --protocol "$PROTOCOL" \ | ||
| 99 | + --num_designs "$NUM_DESIGNS" \ | ||
| 100 | + --budget "$BUDGET" \ | ||
| 101 | + --devices "$DEVICES" \ | ||
| 102 | + --num_workers "$NUM_WORKERS" | ||
| 103 | + | ||
| 104 | +# 捕获执行结果 | ||
| 105 | +EXIT_CODE=$? | ||
| 106 | + | ||
| 107 | +# 计算运行时间 | ||
| 108 | +END_TIME=$(date +%s) | ||
| 109 | +DURATION=$((END_TIME - START_TIME)) | ||
| 110 | + | ||
| 111 | +# ============================================ | ||
| 112 | +# 4. 结果处理 | ||
| 113 | +# ============================================ | ||
| 114 | + | ||
| 115 | +if [ $EXIT_CODE -eq 0 ]; then | ||
| 116 | + print_info "推理执行成功！" | ||
| 117 | + print_info "结束时间: $(date '+%Y-%m-%d %H:%M:%S')" | ||
| 118 | + print_info "总耗时: $((DURATION / 60)) 分 $((DURATION % 60)) 秒" | ||
| 119 | + print_info "输出目录: $WORK_DIR/$OUTPUT_DIR" | ||
| 120 | + | ||
| 121 | + # 显示输出目录内容 | ||
| 122 | + if [ -d "$OUTPUT_DIR" ]; then | ||
| 123 | + print_info "输出目录内容：" | ||
| 124 | + ls -la "$OUTPUT_DIR" 2>/dev/null || print_warning "无法列出输出目录内容" | ||
| 125 | + fi | ||
| 126 | +else | ||
| 127 | + print_error "推理执行失败，退出码: $EXIT_CODE" | ||
| 128 | + print_error "结束时间: $(date '+%Y-%m-%d %H:%M:%S')" | ||
| 129 | + exit $EXIT_CODE | ||
| 130 | +fi | ||
| 131 | + | ||
| 132 | +print_info "脚本执行完毕！" | ||