#!/bin/bash

gpu_idx=0

# # 2/255 
train_eps=0.007843
test_eps=0.007843
# # # SABR
# # eps_shrinkage=0.1
# # relu_shrinkage=0.8
# # L1=1e-6
# # # TAPS
block_sizes="4 17"
L1=2e-6
TAPS_grad_scale=5
num_steps=1
restarts=3
# STAPS
# eps_shrinkage=0.1
# relu_shrinkage=0.8
# L1=2e-6
# block_sizes="17 4"
# TAPS_grad_scale=2

# # 8/255 
# train_eps=0.03137
# test_eps=0.03137
# # # SABR
# # eps_shrinkage=0.7
# # relu_shrinkage=None
# # L1=0
# # # TAPS
# block_sizes="17 4"
# L1=2e-6
# TAPS_grad_scale=5
# # # STAPS
# # eps_shrinkage=0.7
# # relu_shrinkage=None
# # block_sizes="17 4"
# # L1=2e-6
# # TAPS_grad_scale=5

# STAPS Training
# CUDA_VISIBLE_DEVICES=$gpu_idx python mix_train.py \
#        --dataset cifar10 --init fast --fast-reg 0.5 --use-TAPS-training --use-small-box --eps-shrinkage $eps_shrinkage --relu-shrinkage $relu_shrinkage --block-sizes $block_sizes --L1 $L1 --TAPS-grad-scale $TAPS_grad_scale --TAPS-anneal-length 1 \
#        --net cnn_7layer_bn \
#        --lr 0.0005 --L2-reg 0 --lr-milestones 120 140 \
#        --train-eps $train_eps --test-eps $test_eps  --train-steps 20 --test-steps 20  \
#        --train-batch 128 --test-batch 128  \
#        --robust-weight-start 1 --robust-weight-end 1 \
#        --start-epoch-robust-weight 0 --end-epoch-robust-weight 0 \
#        --grad-clip 10 \
#        --n-epochs 160 --start-epoch-eps 1 --end-epoch-eps 81 \
#        --save-dir ./test_models/

# # # TAPS Training
CUDA_VISIBLE_DEVICES=$gpu_idx python mix_train.py \
       --dataset cifar10 --init fast --fast-reg 0.5 --use-TAPS-training --block-sizes $block_sizes --L1 $L1 --TAPS-grad-scale $TAPS_grad_scale --TAPS-anneal-length 10 \
       --net cnn_7layer_bn \
       --lr 0.0005 --L2-reg 0 --lr-milestones 120 140 \
       --train-eps $train_eps --test-eps $test_eps  --train-steps $num_steps --test-steps $num_steps --restarts $restarts \
       --train-batch 128 --test-batch 128  \
       --robust-weight-start 1 --robust-weight-end 1 \
       --start-epoch-robust-weight 0 --end-epoch-robust-weight 0 \
       --grad-clip 10 \
       --n-epochs 160 --start-epoch-eps 1 --end-epoch-eps 81 \
       --save-dir ./test_models/

# # SABR Training
# CUDA_VISIBLE_DEVICES=$gpu_idx python mix_train.py \
#        --dataset cifar10 --init fast --fast-reg 0.5 --use-vanilla-ibp --use-small-box --eps-shrinkage $eps_shrinkage --relu-shrinkage $relu_shrinkage \
#        --net cnn_7layer_bn \
#        --lr 0.0005 --L2-reg 0 --lr-milestones 120 140 \
#        --train-eps $train_eps --test-eps $test_eps  --train-steps 20 --test-steps 20  \
#        --train-batch 128 --test-batch 128  \
#        --robust-weight-start 1 --robust-weight-end 1 \
#        --start-epoch-robust-weight 0 --end-epoch-robust-weight 0 \
#        --grad-clip 10 \
#        --n-epochs 160 --start-epoch-eps 1 --end-epoch-eps 81 \
#        --save-dir ./test_models/

# # IBP Training
# CUDA_VISIBLE_DEVICES=$gpu_idx python mix_train.py \
#        --dataset cifar10 --init fast --fast-reg 0.5 --use-vanilla-ibp \
#        --net cnn_7layer_bn \
#        --lr 0.0005 --L2-reg 0 --lr-milestones 120 140 \
#        --train-eps $train_eps --test-eps $test_eps  --train-steps 20 --test-steps 20  \
#        --train-batch 128 --test-batch 128  \
#        --robust-weight-start 1 --robust-weight-end 1 \
#        --start-epoch-robust-weight 0 --end-epoch-robust-weight 0 \
#        --grad-clip 10 \
#        --n-epochs 160 --start-epoch-eps 1 --end-epoch-eps 81 \
#        --save-dir ./test_models/
