Add sampling

Solor-pikachu · Mar 15, 2024 · afbc59e · afbc59e
1 parent c34e097
commit afbc59e
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 524 deletions.
diff --git a/nila_sampling/.ipynb_checkpoints/Recon-checkpoint.ipynb b/nila_sampling/.ipynb_checkpoints/Recon-checkpoint.ipynb
diff --git a/nila_sampling/.ipynb_checkpoints/mrfft-checkpoint.py b/nila_sampling/.ipynb_checkpoints/mrfft-checkpoint.py
diff --git a/nila_training/.ipynb_checkpoints/train-checkpoint.sh b/nila_training/.ipynb_checkpoints/train-checkpoint.sh
@@ -0,0 +1,5 @@
+MODEL_FLAGS="--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads 4 --learn_sigma True"
+DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear"
+TRAIN_FLAGS="--lr 1e-4 --batch_size 8"
+
+mpiexec -n 8 python scripts/image_train.py --data_dir data/fastmri_complex_value/ --log_dir log $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS
diff --git a/nila_training/scripts/.ipynb_checkpoints/image_train-checkpoint.py b/nila_training/scripts/.ipynb_checkpoints/image_train-checkpoint.py
@@ -26,7 +26,7 @@ def main():
 
     logger.log("creating model and diffusion...")
     print(args_to_dict(args, model_and_diffusion_defaults().keys()))
-    model, diffusion = create_DIT_model_and_diffusion(
+    model, diffusion = create_model_and_diffusion(
         **args_to_dict(args, model_and_diffusion_defaults().keys())
     )
     model.to(dist_util.dev())

diff --git a/nila_training/scripts/image_train.py b/nila_training/scripts/image_train.py
@@ -26,7 +26,7 @@ def main():
 
     logger.log("creating model and diffusion...")
     print(args_to_dict(args, model_and_diffusion_defaults().keys()))
-    model, diffusion = create_DIT_model_and_diffusion(
+    model, diffusion = create_model_and_diffusion(
         **args_to_dict(args, model_and_diffusion_defaults().keys())
     )
     model.to(dist_util.dev())

diff --git a/nila_training/train.sh b/nila_training/train.sh
@@ -2,4 +2,4 @@ MODEL_FLAGS="--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads
 DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear"
 TRAIN_FLAGS="--lr 1e-4 --batch_size 8"
 
-mpiexec -n 8 python scripts/image_train.py --data_dir data/fastmri_complex_value/ --log_dir log_DIT $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS
+mpiexec -n 8 python scripts/image_train.py --data_dir data/fastmri_complex_value/ --log_dir log $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS