release training code

README.md

@@ -8,11 +8,21 @@
 [![Dataset](https://img.shields.io/badge/Dataset-fMRI_Objaverse-3aa0a5.svg?logo=Huggingface)](https://huggingface.co/datasets/Fudan-fMRI/fMRI-Objaverse)
 
 
-
-
 # Introduction
 <img src='./imgs/teaser.png' width="100%">
 
+# Project Updates
+- 🔥 **News**: We have released the training code!
+- 🔥 **Note**: fMRI-Objaverse: https://huggingface.co/datasets/Fudan-fMRI/fMRI-Objaverse
+- 🔥 **Note**: fMRI-Shape: https://huggingface.co/datasets/Fudan-fMRI/fMRI-Shape
+
+# Environment Setup
+
+```bash
+git clone https://github.com/JianxGao/MinD-3D.git
+cd MinD-3D
+bash env_install.sh
+```
 
 # Dataset and checkpoints
 You can download fMRI-Shape by this link: https://huggingface.co/datasets/Fudan-fMRI/fMRI-Shape.
@@ -23,14 +33,28 @@ You can download the weight of subject 1 through the link:
 
 https://drive.google.com/file/d/1ni4g1iCvdpoi2xYtmydr_w3XA5PpNrvm/view?usp=sharing
 
-# Environment Setup
+
+
+# Train
 
 ```bash
-git clone https://github.com/JianxGao/MinD-3D.git
-cd MinD-3D
-bash env_install.sh
+CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port=25645 \
+ train_stage1.py --sub_id 0001 --ddp \
+ --config ./configs/mind3d.yaml \
+ --out_dir sub01_stage1 --batchsize 8
 ```
 
+```bash
+CUDA_VISIBLE_DEVICES=1 python -m torch.distributed.launch --nproc_per_node=1 --master_port=25645 \
+ train_stage2.py --sub_id 0001 --ddp \
+ --config ./configs/mind3d.yaml \
+ --out_dir sub01_stage2 --batchsize 2
+```
+
+You can access the quantized features for training through the link: https://drive.google.com/file/d/1R8IpG1bligLAfHkLQ2COrfTIkay14AEm/view?usp=drive_link.
+
+
+
 # Inference
 
 ```bash

configs/mind3d.yaml

@@ -101,9 +101,6 @@ training:
   warmup_iters: 150
   max_iters: 3000
 
-  n_workers: 8
-  n_workers_val: 4
-
 test:
   threshold: 0.2
   eval_mesh: true

dataset/test_list.txt

@@ -0,0 +1,104 @@
+02691156/ed35478403ae873943cf31d2bcc8f4.npy
+02828884/7aca4c47c6861f2122445e799be0f18.npy
+02958343/aebd98c5d7e8150b709ce7955adef61b.npy
+02828884/51197dd964f598fdcb2a965e75be701c.npy
+04530566/e5524ebe2d9f1e8c47caa6b5f3c46e10.npy
+04530566/51f269a4ff48c2214de203ef6842ee61.npy
+02933112/c7165635f2288945585ed17f54616d23.npy
+02958343/7f4dbefeaeda54628ccbb4eb8965b05.npy
+04401088/725cf43af935ee023832db239383a26e.npy
+04256520/c0b61d423f449b3c6f37a8be59c15bbb.npy
+02958343/99efa2111f3be22d5fcb05674109534a.npy
+04256520/a86bf379ca2ba8d8f2d5a0d235580548.npy
+02958343/68b4a03cd8d49553e98e65cfc071e8a0.npy
+03636649/78a11c0b8e964c9b41657e31b569b105.npy
+04090263/12b346d7a60328ba9be47ac7c7e37815.npy
+02933112/8d3d624b37b4803c33110e4a1f22f6dc.npy
+04530566/c680824e615c772d18d5c097fe7ed300.npy
+04401088/73b87e2025cfdb6daf19fa8e8218bf64.npy
+04256520/6e4fab35f29925f14f438613584cbdf7.npy
+04379243/798a07e42d76013582695d8aaeacccc5.npy
+04379243/40a402e1d949364a104ceb84075e40d6.npy
+03211117/5b2585db758fac87bed75a82d9bfc7c.npy
+03001627/b526dccbbe8417bd6cf45f1c1a3d29af.npy
+02828884/23378697712a8dc2d2b12aa6a0f050b3.npy
+04401088/f6679591247bc328afb07a946d621b3c.npy
+04530566/e23e4845d3a990d25fc186cfb5f359d1.npy
+02691156/eefb4a1fcc2ca4d8894e22af7ae821f8.npy
+03211117/34808bb66b2bfa07ab7f9d2da6fc61cf.npy
+04090263/c77af288e7c373979905a3939c890ad2.npy
+03211117/7a2d9d3d185cff794b85c79082892df6.npy
+04379243/45a09b1ce3111e4f22f4fabdf1ee0670.npy
+02933112/3c28e4a1a5154ecc59350d819542ec7.npy
+02933112/c365225746f2eb94ce5d6e9371bb5c33.npy
+04401088/69e45e401c1d05b767d8789d3a6519d0.npy
+02691156/e8e1b765fdf5edfa14c19f41d007670e.npy
+02691156/cbb5ed444249f0a9e90916c8815b5c43.npy
+04379243/c876b151a22b59a24702663c3b3faf8f.npy
+02691156/61c804d4751079962df73a6189a3b87.npy
+04401088/e85bac837e951f3780ed245d94a6a268.npy
+04090263/42b7e8c381cb392740112771efe22586.npy
+03636649/42129d6b04cb9038b2f7093e7efce142.npy
+03691459/77c70fcd8ebcfbd4c83db093c68b1d7b.npy
+04256520/cba1446e98640f603ffc853fc4b95a17.npy
+03001627/ba1adc474cfa2f429ab926a7e944e795.npy
+03001627/97df0e7773e51feb331fc18393f04d2a.npy
+03211117/5ebad73554fd866830bcd442adebae0f.npy
+04090263/dcc4bb465d403801e1f2a1daf140ac9f.npy
+04379243/eb8a34ab2eb9a26fe93eff560d16430.npy
+04379243/20bb1a34891588c1cce90f0d61ed7a70.npy
+03636649/4dc9f86110ea40aaa570c6c691c987a8.npy
+03636649/2d638c6b6b2feb9248da169d95204ce2.npy
+04090263/abcfbaaa409d62001d263b1fb57ddbcf.npy
+02828884/3ff4401d898b488a601104cd2d998272.npy
+02933112/5bfc4cb92e9d3b1f79b0363dd1368699.npy
+02958343/8df9feeed8678aa6bdc56c6b6d74b51a.npy
+04401088/9dc0c2ca8de18e28b4c18e3209e12990.npy
+03001627/d75d41af68f4bf7f9af4053b8815b239.npy
+03636649/bb3c2fba4e1b712357ddce8e42aa6e90.npy
+04530566/7e6aa3337c2e52ff196f35b142096073.npy
+02691156/d6cc7ffd165de05e2e2de556c136fbc8.npy
+03636649/cd80cc92cf732e8d8a17805dbfb751e2.npy
+04256520/d13a2ccdbb7740ea83a0857b7b9398b1.npy
+03636649/41fce4197aad78f76fab5a5e55726ee7.npy
+04256520/3b7f2867c3d56290267536df1b0cc6.npy
+02828884/ea6da848da3f603f81faebbdea6bd9be.npy
+02828884/8e260916773769032b7072d8dd8cf43d.npy
+02933112/534d0386ab0bb650ae96150e4aa362f9.npy
+04401088/457b4310870bcf047f4e0c14aca1a926.npy
+02958343/cf698011f90ac05f253c03b7df20edd5.npy
+02933112/2375dfeeab7fd8678430b7446f9e1252.npy
+04530566/e853f137948055b188d5590328ce0ddf.npy
+04379243/eb62d577d4de8bcbdda7f1d27cdf247c.npy
+03691459/d97a4771b9d796a22d1a40e420084a6.npy
+04090263/421da608bba18bf2688fc2fc0a14edd5.npy
+02958343/2f0ca88dc76e39a7ff1370452e3a0154.npy
+02933112/c7c782fd5976ed7214038d588fd1342f.npy
+03001627/7b405c1d6d2dbea9f91663a74ccd2338.npy
+02958343/5695a98b66b2b9a9c81ddeca50aa3117.npy
+04090263/6e91124a4ecaa1a365564d9c2ed3b493.npy
+04090263/d4312787c89edb399a6e43b878d5b335.npy
+04401088/17072ff9ec47bfadc3fd2392bc923b2f.npy
+03211117/b7a259bd391ad43c69b642ba22680f87.npy
+03211117/52e60434725b8bf8f0975d5dbdd766e.npy
+03691459/2eacd8302d46093ee138e561e8f9c254.npy
+04530566/a51410ca51a40c2529412f6cb23454eb.npy
+02691156/c12362906beecc796059aeb39e538eb2.npy
+04256520/312ee9c8bc6ecec5ce59ea70152320fa.npy
+02691156/b5d0ae4f723bce81f119374ee5d5f944.npy
+03001627/7f9733a6370440dcc8687ff9b0b4e4ac.npy
+03691459/97bf4aac2d956c1d5d9e9a1d5cade7db.npy
+03001627/5bd08666f29e946c988a5637d6d8f36a.npy
+03691459/8f920f26ead4d592cf66d96a8426fc6.npy
+03691459/336fcd767dadec9fb611871e114769d9.npy
+04530566/31f7c298ea24bb99847c4e55bb965ab0.npy
+04256520/cc7b690e4d86b471397aad305ec14786.npy
+03636649/6d5d4e506d6630b193528e10aeca0aae.npy
+02828884/969f830a531682cef155d75bbf62b80.npy
+03211117/4a21927379f965a9e4b68d3b17c43658.npy
+04379243/bf9df04aea08b0f0f51f77a6d7299806.npy
+02828884/1b77f2fbf17f28e726cd2499b0c05f5a.npy
+03211117/72b55fd0ec85d68164ac9dbe3ce7e6be.npy
+03691459/3dc77d756da1ff5eb3b44ad99a214777.npy
+03001627/6455bf241fcbe3a03adc3470b30138f3.npy
+03691459/774214a85b87204c33dbc1ea1716bde.npy