File size: 2,115 Bytes
a2ffd07 | 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 | #!/bin/bash
# Define the parameters
model_name="Salesforce/blip-image-captioning-base"
tok_name="Salesforce/blip-image-captioning-base"
batch_size=16
max_tokens=5000000000000
sae_name="batchtopk"
lr=1e-4
expansion_factor=8
k=8
auxk=128
auxk_coef=0.03125
device_id=1
max_epochs=3
dead_tokens_threshold=1000000
log_every_n_steps=100
save_every_n_training_steps=100
save_step=false
use_loss_var=true
num_workers=4
hf_dataset="pixparse/cc3m-wds"
dtype="float16"
# change local_val_path
local_train_path="CC3M-Dataset/preproc/CC3M-Dataset/cc3m_images/train"
local_val_path="CC3M-Dataset/preproc/CC3M-Dataset/cc3m_images/val"
wandb_key=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
cd ..
# Loop over 12 layers
for layer in {10..10}
do
# hook names: vision_model.encoder.layers.${layer}.hook_resid_pre,
# text_decoder.bert.encoder.layer.${layer}.attention.self.hook_resid_pre,
# text_decoder.bert.encoder.layer.${layer}.crossattention.self.hook_resid_pre,
hook_name="vision_model.encoder.layers.${layer}.hook_resid_post"
wandb login --relogin $wandb_key
wandb login
python Train_SAE_Blip.py \
--model_name "$model_name" \
--tok_name "$tok_name" \
--layer "$layer" \
--batch_size "$batch_size" \
--hook_name "$hook_name" \
--sae_name "$sae_name" \
--lr "$lr" \
--expansion_factor "$expansion_factor" \
--k "$k" \
--auxk "$auxk" \
--max_epochs "$max_epochs" \
--dead_tokens_threshold "$dead_tokens_threshold" \
--log_every_n_steps "$log_every_n_steps" \
--save_every_n_training_steps "$save_every_n_training_steps" \
--save_step "$save_step" \
--use_loss_var "$use_loss_var" \
--max_tokens "$max_tokens" \
--auxk_coef "$auxk_coef" \
--num_workers "$num_workers" \
--hf_dataset "$hf_dataset" \
--local_train_path "$local_train_path" \
--local_val_path "$local_val_path" \
--device_id "$device_id" \
--dtype "$dtype" \
done |