From eb22dd3916890c7be115ec0dc0fb82fd64668f8d Mon Sep 17 00:00:00 2001
From: ghliu <ghliu@gatech.edu>
Date: Fri, 1 Dec 2023 17:48:57 -0500
Subject: [PATCH] add scripts

---
 .gitignore                   |  5 +++++
 scripts/download_img_data.sh | 27 +++++++++++++++++++++++++++
 scripts/train_constr_gen.sh  | 21 +++++++++++++++++++++
 scripts/train_img.sh         | 20 ++++++++++++++++++++
 4 files changed, 73 insertions(+)
 create mode 100644 scripts/download_img_data.sh
 create mode 100644 scripts/train_constr_gen.sh
 create mode 100644 scripts/train_img.sh

diff --git a/.gitignore b/.gitignore
index 68bc17f..464bb57 100644
--- a/.gitignore
+++ b/.gitignore
@@ -158,3 +158,8 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+#
+data/*
+results/*
+.log/*
diff --git a/scripts/download_img_data.sh b/scripts/download_img_data.sh
new file mode 100644
index 0000000..4485c49
--- /dev/null
+++ b/scripts/download_img_data.sh
@@ -0,0 +1,27 @@
+
+DATASET=$1
+
+echo "Note: available dataset are ffhq, afhqv2 "
+echo "Specified [$DATASET]"
+
+if [ "$DATASET" == afhqv2 ]; then
+    # afhqv2-64x64.zip
+    gdown -O data/ 17mH2b3oQ1CjqNshDBqZWapFoZbl2Rcpy
+
+    # 3x64x64 orthobasis
+    gdown -O data/ 1JYm8O46Y2bZOySDghyODp0ud-5lZcKnc
+
+    # network (VP, uncond)
+    gdown -O data/ 1A1upwcOmLfcHV0j4n2Fa8rtOnSVcotmh
+fi
+
+if [ "$DATASET" == ffhq ]; then
+    # ffhq-64x64.zip
+    gdown -O data/ 1NQ40CmBxFiAkd2SoWUTbhmm3VAyWVerF
+
+    # 3x64x64 orthobasis
+    gdown -O data/ 1JYm8O46Y2bZOySDghyODp0ud-5lZcKnc
+
+    # network (VP, uncond)
+    gdown -O data/ 1nz2M3gn-CbQ2z2_oKRSA5M_VwQzrtXdt
+fi
diff --git a/scripts/train_constr_gen.sh b/scripts/train_constr_gen.sh
new file mode 100644
index 0000000..accf699
--- /dev/null
+++ b/scripts/train_constr_gen.sh
@@ -0,0 +1,21 @@
+# ball
+python train_constr_gen.py --constraint ball --xdim 2 --p0 moon    --name moon
+python train_constr_gen.py --constraint ball --xdim 2 --p0 spiral  --name spiral
+python train_constr_gen.py --constraint ball --xdim 2 --p0 gmm     --name gmm
+python train_constr_gen.py --constraint ball --xdim 6 --p0 gmm_nd  --name gmm6d
+python train_constr_gen.py --constraint ball --xdim 8 --p0 gmm_nd  --name gmm8d
+python train_constr_gen.py --constraint ball --xdim 20 --p0 gmm_nd --name gmm20d --num-itr 100000
+
+# simplex
+python train_constr_gen.py --constraint simplex --xdim 2  --p0 dirichlet1   --name simplex3d
+python train_constr_gen.py --constraint simplex --xdim 2  --p0 dirichlet2   --name simplex3dv2
+python train_constr_gen.py --constraint simplex --xdim 6  --p0 dirichlet3   --name simplex7d
+python train_constr_gen.py --constraint simplex --xdim 8  --p0 dirichlet4   --name simplex9d
+python train_constr_gen.py --constraint simplex --xdim 19 --p0 dirichlet_nd --name simplex20d --num-itr 100000
+
+# cube
+python train_constr_gen.py --constraint cube --xdim 2 --p0 cube --name cube2d
+python train_constr_gen.py --constraint cube --xdim 3 --p0 cube --name cube3d
+python train_constr_gen.py --constraint cube --xdim 6 --p0 cube --name cube6d
+python train_constr_gen.py --constraint cube --xdim 8 --p0 cube --name cube8d
+python train_constr_gen.py --constraint cube --xdim 20 --p0 cube --name cube20d --num-itr 100000
diff --git a/scripts/train_img.sh b/scripts/train_img.sh
new file mode 100644
index 0000000..028c1ee
--- /dev/null
+++ b/scripts/train_img.sh
@@ -0,0 +1,20 @@
+DATASET=$1
+
+echo "Note: available dataset are ffhq, afhqv2 "
+echo "Specified [$DATASET]"
+
+# these arguments should NOT be modified unless pretrain models change
+AFHQv2="--pretrain=data/edm-afhqv2-64x64-uncond-vp.pt --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15"
+FFHQ="  --pretrain=data/edm-ffhq-64x64-uncond-vp.pt   --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15"
+
+# `batch-gpu` is the maximum batch per gpu. **Current values are tuend on TITAN RTX**
+# `snap` controls the frequency (in kimg) of saving network snapshot
+# `dump` controls the frequency (in kimg) of dumpping training stats
+
+if [ "$DATASET" == afhqv2 ]; then
+    python train_watermark.py --data=data/afhqv2-64x64.zip --outdir=results $AFHQv2 --batch-gpu=64 --snap=20 --dump=20
+fi
+
+if [ "$DATASET" == ffhq ]; then
+    python train_watermark.py --data=data/ffhq-64x64.zip --outdir=results $FFHQ   --batch-gpu=64 --snap=10 --dump=10
+fi