jsflow / REG /train.sh
xiangzai's picture
Add files using upload-large-folder tool
b65e56d verified
#!/usr/bin/env bash
# REG/train.py:与主仓库类似,可单独指定数据根目录与预处理 cls 特征目录。
# 数据布局:${DATA_DIR}/imagenet_256_vae/ 下 VAE latent;
# ${SEMANTIC_FEATURES_DIR}/ 下 img-feature-*.npy + dataset.json(与 parallel_encode 一致)。
NUM_GPUS=4
# ------------ 按本机路径修改 ------------
DATA_DIR="/gemini/space/zhaozy/dataset/Imagenet/imagenet_256"
SEMANTIC_FEATURES_DIR="/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0"
# 后台示例(与主实验脚本风格一致):
# nohup bash train.sh > jsflow-experiment.log 2>&1 &
nohup accelerate launch --multi_gpu --num_processes "${NUM_GPUS}" --mixed_precision bf16 train.py \
--report-to wandb \
--allow-tf32 \
--mixed-precision bf16 \
--seed 0 \
--path-type linear \
--prediction v \
--weighting uniform \
--model SiT-XL/2 \
--enc-type dinov2-vit-b \
--encoder-depth 8 \
--proj-coeff 0.5 \
--output-dir exps \
--exp-name jsflow-experiment-0.75-0.01 \
--batch-size 256 \
--data-dir "${DATA_DIR}" \
--semantic-features-dir "${SEMANTIC_FEATURES_DIR}" \
--learning-rate 0.00005 \
--t-c 0.75 \
--cls 0.01 \
--ot-cls \
> jsflow-experiment.log 2>&1 &
# 说明:
# - 不使用预处理特征、改在线抽 DINO 时:去掉 --semantic-features-dir,并保证 data-dir 为 REG 原布局
# (imagenet_256_vae + vae-sd)。
# - 关闭 minibatch OT:追加 --no-ot-cls。
# - 主仓库 train.py 中的 --weight-ratio / --semantic-reg-coeff / --repa-* 等为本 REG 脚本未实现项;
# 投影强度请用 --proj-coeff,cls 流损失权重用 --cls。