CaptchaBench/install_new_envs.sh at main · Linzy19/CaptchaBench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
#!/usr/bin/env bash
# ============================================================
# install_new_envs.sh — 创建新增 2 个攻击方法的 conda 环境：coa / anyattack
#   与 install_all_envs.sh 同范式（conda create --prefix + pip 安装）。
#
# 依赖推断来源：
#   CoA   (Chain_of_Attack/train.py, run_coa_cluster.py)：
#     torch torchvision、openai-clip(含 ftfy/regex)、transformers(GPT2 -> ClipCap)、
#     numpy、pillow、tqdm。ClipCap 权重 conceptual_weights.pt 已在
#     Chain_of_Attack/clip_prefix_model/（无需 pip）。
#   AnyAttack (AnyAttack/models/model.py, run_anyattack_cluster.py)：
#     torch torchvision、ftfy regex(供自带的 models/clip 使用)、numpy、pillow、tqdm。
#     decoder 权重 coco_cos.pt 已在 AnyAttack/checkpoints/（无需 pip）。
#     注意:AnyAttack 自带 models/clip,不要再 pip 装 openai-clip(避免命名冲突)。
#
# 用法：
#   bash install_new_envs.sh                 # 创建两个环境（已存在则跳过）
#   bash install_new_envs.sh --force         # 删除重建
#   bash install_new_envs.sh --skip_coa      # 只建 anyattack
#   bash install_new_envs.sh --cuda_version 12.1
# ============================================================
set -eo pipefail
trap 'echo "[错误] 脚本在第 ${LINENO} 行失败，退出码：$?" >&2' ERR

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ---- 配置（与 install_all_envs.sh 对齐）----
FORCE_RECREATE=0
PYTORCH_CUDA_VERSION="11.8"
PYTHON_VERSION="3.10"
CONDA_ENV_DIR="/apdcephfs_qy3/share_470749/lzy_private/conda_env_lzy"
SKIP_COA=0
SKIP_ANYATTACK=0

while [[ $# -gt 0 ]]; do
    case "$1" in
        --force)          FORCE_RECREATE=1; shift ;;
        --skip_coa)       SKIP_COA=1; shift ;;
        --skip_anyattack) SKIP_ANYATTACK=1; shift ;;
        --cuda_version)   PYTORCH_CUDA_VERSION="$2"; shift 2 ;;
        --python_version) PYTHON_VERSION="$2"; shift 2 ;;
        *) echo "[错误] 未知参数：$1" >&2; exit 1 ;;
    esac
done

# ---- conda 初始化 ----
conda_base="$(conda info --base 2>/dev/null)" || { echo "[错误] 未找到 conda" >&2; exit 1; }
# shellcheck source=/dev/null
source "${conda_base}/etc/profile.d/conda.sh"

create_or_skip_env() {
    local env_name="$1" python_ver="$2"
    local env_path="${CONDA_ENV_DIR}/${env_name}"
    if [[ -d "${env_path}" && -f "${env_path}/bin/python" ]]; then
        if [[ ${FORCE_RECREATE} -eq 1 ]]; then
            echo "[重建] 删除并重建：${env_path}"; rm -rf "${env_path}"
        else
            echo "[跳过] 环境已存在：${env_path}（--force 可强制重建）"; return 0
        fi
    fi
    echo "[创建] ${env_path} (Python ${python_ver})"
    conda create --prefix "${env_path}" python="${python_ver}" -y
}

install_pytorch() {
    local cuda_ver="$1"
    local cuda_suffix; cuda_suffix="cu$(echo "${cuda_ver}" | tr -d '.')"
    echo "[安装] PyTorch (CUDA ${cuda_ver})..."
    pip install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/${cuda_suffix}"
}

# ============================================================
# coa 环境
# ============================================================
install_coa() {
    echo ""; echo "=== [安装] 环境：coa (Chain-of-Attack) ==="
    create_or_skip_env "coa" "${PYTHON_VERSION}"
    set +e; conda activate "${CONDA_ENV_DIR}/coa"; set -e
    install_pytorch "${PYTORCH_CUDA_VERSION}"
    # OpenAI CLIP（提供 import clip）+ transformers（GPT2 -> ClipCap）
    echo "[安装] coa 依赖（需要网络）..."
    pip install ftfy regex tqdm numpy pillow
    pip install "transformers>=4.30,<5"   # GPT2Tokenizer/GPT2LMHeadModel；<5 避免 AdamW 移除问题
    pip install git+https://github.com/openai/CLIP.git
    echo "[验证] coa 环境..."
    python -c "import torch, clip, transformers; print('coa OK, CUDA:', torch.cuda.is_available(), '| transformers', transformers.__version__)"
}

# ============================================================
# anyattack 环境
# ============================================================
install_anyattack() {
    echo ""; echo "=== [安装] 环境：anyattack ==="
    create_or_skip_env "anyattack" "${PYTHON_VERSION}"
    set +e; conda activate "${CONDA_ENV_DIR}/anyattack"; set -e
    install_pytorch "${PYTORCH_CUDA_VERSION}"
    # AnyAttack 自带 models/clip，只需其依赖 ftfy/regex；不装 openai-clip
    echo "[安装] anyattack 依赖（需要网络）..."
    pip install ftfy regex tqdm numpy pillow
    echo "[验证] anyattack 环境..."
    cd "${SCRIPT_DIR}/AnyAttack"
    python -c "import torch; import models.clip as clip; from models.model import CLIPEncoder, Decoder; print('anyattack OK, CUDA:', torch.cuda.is_available())"
}

[[ ${SKIP_COA} -eq 0 ]]       && install_coa
[[ ${SKIP_ANYATTACK} -eq 0 ]] && install_anyattack

echo ""; echo "[完成] 新环境安装结束。"
echo "提示：若集群无外网，需提前缓存 CLIP ViT-B/32（放 .clip_cache 并 export CLIP_DOWNLOAD_ROOT）"
echo "      以及 HuggingFace gpt2（设置 HF_HOME 指向已缓存目录）。"