Fix mito_mitoEM_H OOM training

Author: Donglai Wei

justfile

@@ -121,13 +121,13 @@ slurm partition num_cpu num_gpu cmd constraint='' mem='32G':
            --output=slurm_outputs/slurm-%j.out \
            --error=slurm_outputs/slurm-%j.err \
            --nodes=1 \
-           --ntasks={{num_gpu}} \
+           --ntasks=1 \
            --gpus-per-node={{num_gpu}} \
            --cpus-per-task={{num_cpu}} \
            --mem={{mem}} \
            --time=$time_limit \
            $constraint_flag \
-           --wrap="mkdir -p \$HOME/.just && export JUST_TEMPDIR=\$HOME/.just TMPDIR=\$HOME/.just && source /projects/weilab/weidf/lib/miniconda3/bin/activate pytc && cd $PWD && srun --ntasks={{num_gpu}} --ntasks-per-node={{num_gpu}} {{cmd}}"
+           --wrap="mkdir -p \$HOME/.just && export JUST_TEMPDIR=\$HOME/.just TMPDIR=\$HOME/.just NCCL_SOCKET_FAMILY=AF_INET && source /projects/weilab/weidf/lib/miniconda3/bin/activate pytc && cd $PWD && srun --ntasks=1 {{cmd}}"
 
 # Alias for slurm (kept for backward compatibility)
 slurm-sh partition num_cpu num_gpu cmd constraint='' mem='32G':

tutorials/mito_mitoEM_H.yaml

@@ -1,7 +1,28 @@
 _base_: mito_mitoEM_common.yaml
 
 experiment_name: mitoem30h_mednext_sdt
+system:
+  training:
+    num_workers: 2
+    batch_size: 2
+  inference:
+    batch_size: 1
+model:
+  input_size:
+  - 16
+  - 256
+  - 256
+  output_size:
+  - 16
+  - 256
+  - 256
 data:
+  use_preloaded_cache: false
+  persistent_workers: false
+  patch_size:
+  - 16
+  - 256
+  - 256
   train_image:
   - EM30-H/im_train.h5
   train_label:
@@ -10,6 +31,15 @@ data:
   - EM30-H/im_val.h5
   val_label:
   - EM30-H/mito_val-v2.h5
+optimization:
+  accumulate_grad_batches: 4
+inference:
+  sliding_window:
+    window_size:
+    - 16
+    - 256
+    - 256
+    sw_batch_size: 1
 monitor:
   checkpoint:
     dirpath: outputs/mitoem30h_mednext_sdt/checkpoints/