diff --git a/.github/issue_stale.yaml b/.github/issue_stale.yaml
deleted file mode 100644
index 17948d3..0000000
--- a/.github/issue_stale.yaml
+++ /dev/null
@@ -1,16 +0,0 @@
-# Number of days of inactivity before an issue becomes stale
-daysUntilStale: 7
-# Number of days of inactivity before a stale issue is closed
-daysUntilClose: 7
-# Issues with these labels will never be considered stale
-exemptLabels:
- - backlog
-# Label to use when marking an issue as stale
-staleLabel: stale
-# Comment to post when marking an issue as stale. Set to `false` to disable
-markComment: >
- This issue has been automatically marked as stale because it has not had
- recent activity. It will be closed if no further activity occurs. Thank you
- for your contributions.
-# Comment to post when closing a stale issue. Set to `false` to disable
-closeComment: false
\ No newline at end of file
diff --git a/.github/workflows/issue_stale.yaml b/.github/workflows/issue_stale.yaml
new file mode 100644
index 0000000..7c11c2b
--- /dev/null
+++ b/.github/workflows/issue_stale.yaml
@@ -0,0 +1,30 @@
+name: Close inactive issues
+on:
+ schedule:
+ - cron: "35 11 * * 5"
+
+env:
+ DAYS_BEFORE_ISSUE_STALE: 30
+ DAYS_BEFORE_ISSUE_CLOSE: 14
+
+jobs:
+ close-issues:
+ runs-on: ubuntu-latest
+ permissions:
+ issues: write
+ pull-requests: write
+ steps:
+ - uses: actions/stale@v5
+ with:
+ days-before-issue-stale: ${{ env.DAYS_BEFORE_ISSUE_STALE }}
+ days-before-issue-close: ${{ env.DAYS_BEFORE_ISSUE_CLOSE }}
+ stale-issue-label: "stale"
+ stale-issue-message: |
+ This issue is stale because it has been open for ${{ env.DAYS_BEFORE_ISSUE_STALE }} days with no activity.
+ It will be closed if no further activity occurs. Let us know if you still need help!
+ close-issue-message: |
+ This issue is being closed because it has been stale for ${{ env.DAYS_BEFORE_ISSUE_CLOSE }} days with no activity.
+ If you still need help, please feel free to leave comments.
+ days-before-pr-stale: -1
+ days-before-pr-close: -1
+ repo-token: ${{ secrets.GITHUB_TOKEN }}
\ No newline at end of file
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 0000000..c6f98a3
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,43 @@
+# Contributing to OpenSceneFlow
+
+We want to make contributing to this project as easy and transparent as possible. We welcome any contributions, from bug fixes to new features. If you're interested in adding your own scene flow method, this guide will walk you through the process.
+
+## Adding a New Method
+
+Here is a quick guide to integrating a new method into the OpenSceneFlow codebase.
+
+### 1. Data Preparation
+
+All data is expected to be processed into the `.h5` format. Each file represents a scene, and within the file, each data sample is indexed by a `timestamp` key.
+
+For more details on the data processing pipeline, please see the [Data Processing README](./dataprocess/README.md#process).
+
+### 2. Model Implementation
+
+All model source files are located in [`src/models`](./src/models). When adding your model, please remember to import your new model class in the [`src/models/__init__.py`](./src/models/__init__.py) file. Don't forget to add your model conf files in [`conf/model`](./conf/model).
+
+ * **For Feed-Forward Methods:** You can use `deflow` and `fastflow3d` as implementation examples.
+ * **For Optimization-Based Methods:** Please refer to `nsfp` and `fastnsf` for guidance on structure and integration. A detailed example can be found in the [NSFP model file](./src/models/nsfp.py).
+
+### 3. Custom Loss Functions
+
+All loss functions are defined in [`src/lossfuncs.py`](./src/lossfuncs.py). If your model requires a new loss function, you can add it to this file by following the pattern of the existing functions. SeFlow provided a self-supervised loss function example for all feed-forward methods. Feel free to check.
+
+### 4.1 Training a Feed-Forward Model
+
+1. Add a configuration file for your new model in the [`conf/model`](./conf/model) directory.
+2. Begin training by running the following command:
+ ```bash
+ python train.py model=your_model_name
+ ```
+3. **Note:** If your model's output dictionary (`res_dict`) has a different structure from the existing models, you may need to add a new pattern in the `training_step` and `validation_step` methods in the main training script.
+
+### 4.2 Running an Optimization-Based Model
+
+Our framework supports multi-GPU execution for optimization-based methods out of the box. You can follow the structure of existing methods like NSFP to run and evaluate your model.
+
+-----
+
+Once the steps above are completed, other parts of the framework, such as evaluation (`eval`) and visualization (`save`), should integrate with your new model accordingly.
+
+Thank you for your contribution!
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index 1ecc08d..0a50ff6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,14 +1,8 @@
-# check more: https://hub.docker.com/r/nvidia/cuda
-FROM nvidia/cuda:11.7.1-devel-ubuntu20.04
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
ENV DEBIAN_FRONTEND noninteractive
+LABEL maintainer="Qingwen Zhang "
-RUN apt update && apt install -y git curl vim rsync htop
-
-RUN curl -o ~/miniforge3.sh -LO https://github.com/conda-forge/miniforge/releases/latest/download/miniforge3-Linux-x86_64.sh && \
- chmod +x ~/miniforge3.sh && \
- ~/miniforge3.sh -b -p /opt/conda && \
- rm ~/miniforge3.sh && \
- /opt/conda/bin/conda clean -ya && /opt/conda/bin/conda init bash
+RUN apt update && apt install -y git tmux curl vim rsync libgl1 libglib2.0-0 ca-certificates
# install zsh and oh-my-zsh
RUN apt update && apt install -y wget git zsh tmux vim g++
@@ -20,15 +14,23 @@ RUN sh -c "$(wget -O- https://github.com/deluan/zsh-in-docker/releases/download/
-p https://github.com/zsh-users/zsh-syntax-highlighting
RUN printf "y\ny\ny\n\n" | bash -c "$(curl -fsSL https://raw.githubusercontent.com/Kin-Zhang/Kin-Zhang/main/scripts/setup_ohmyzsh.sh)"
-RUN /opt/conda/bin/conda init zsh && /opt/conda/bin/mamba init zsh
+RUN /opt/conda/bin/conda init zsh
# change to conda env
ENV PATH /opt/conda/bin:$PATH
+RUN /opt/conda/bin/conda config --set solver libmamba
-RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/KTH-RPL/OpenSceneFlow.git
+RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/Kin-Zhang/OpenSceneFlow
WORKDIR /home/kin/workspace/OpenSceneFlow
-RUN apt-get update && apt-get install libgl1 -y
+
# need read the gpu device info to compile the cuda extension
-RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f environment.yaml
-# environment for dataprocessing inlucdes data-api
-RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f envsftool.yaml
+RUN /opt/conda/bin/pip install -r /home/kin/workspace/OpenSceneFlow/requirements.txt
+RUN /opt/conda/bin/pip install FastGeodis --no-build-isolation
+RUN /opt/conda/bin/pip install --no-cache-dir -e ./assets/cuda/chamfer3D && /opt/conda/bin/pip install --no-cache-dir -e ./assets/cuda/mmcv
+
+# environment for dataprocessing includes data-api
+RUN /opt/conda/bin/conda env create -f envsftool.yaml
+RUN /opt/conda/envs/sftool/bin/pip install numpy==1.22
+
+# clean up apt cache
+RUN rm -rf /var/lib/apt/lists/* && rm -rf /root/.cache/pip
diff --git a/LICENSE b/LICENSE
index 694508c..07acd9b 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
BSD 3-Clause License
-Copyright (c) 2024, Robotics, Perception and Learning @KTH
+Copyright (c) 2024, Qingwen Zhang, Robotics, Perception and Learning @KTH
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
diff --git a/README.md b/README.md
index ed2af54..51d5307 100644
--- a/README.md
+++ b/README.md
@@ -43,13 +43,13 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer
- [x] [FastFlow3D](https://arxiv.org/abs/2103.01306): RA-L 2021, a basic backbone model.
- [x] [ZeroFlow](https://arxiv.org/abs/2305.10424): ICLR 2024, their pre-trained weight can covert into our format easily through [the script](tools/zerof2ours.py).
-- [ ] [NSFP](https://arxiv.org/abs/2111.01253): NeurIPS 2021, faster 3x than original version because of [our CUDA speed up](assets/cuda/README.md), same (slightly better) performance. Done coding, public after review.
-- [ ] [FastNSF](https://arxiv.org/abs/2304.09121): ICCV 2023. SSL optimization-based. Done coding, public after review.
+- [x] [NSFP](https://arxiv.org/abs/2111.01253): NeurIPS 2021, faster 3x than original version because of [our CUDA speed up](assets/cuda/README.md), same (slightly better) performance.
+- [x] [FastNSF](https://arxiv.org/abs/2304.09121): ICCV 2023. SSL optimization-based.
- [ ] [ICP-Flow](https://arxiv.org/abs/2402.17351): CVPR 2024. SSL optimization-based. Done coding, public after review.
-💡: Want to learn how to add your own network in this structure? Check [Contribute section](assets/README.md#contribute) and know more about the code. Fee free to pull request and your bibtex [here](#cite-us).
+💡: Want to learn how to add your own network in this structure? Check [Contribute section](CONTRIBUTING.md#adding-a-new-method) and know more about the code. Fee free to pull request and your bibtex [here](#cite-us).
---
@@ -76,6 +76,7 @@ cd OpenSceneFlow && mamba env create -f environment.yaml
# You may need export your LD_LIBRARY_PATH with env lib
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/kin/mambaforge/lib
```
+We also provide [requirements.txt](requirements.txt), please check usage through [Dockerfile](Dockerfile).
### Docker (Recommended for Isolation)
@@ -86,11 +87,13 @@ You always can choose [Docker](https://en.wikipedia.org/wiki/Docker_(software))
docker pull zhangkin/opensf
# run container
-docker run -it --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name opensceneflow zhangkin/opensf /bin/zsh
+docker run -it --net=host --gpus all -v /dev/shm:/dev/shm -v /home/kin/data:/home/kin/data --name opensf zhangkin/opensf /bin/zsh
+
# and better to read your own gpu device info to compile the cuda extension again:
+cd /home/kin/workspace/OpenSceneFlow && git pull
cd /home/kin/workspace/OpenSceneFlow/assets/cuda/mmcv && /opt/conda/envs/opensf/bin/python ./setup.py install
cd /home/kin/workspace/OpenSceneFlow/assets/cuda/chamfer3D && /opt/conda/envs/opensf/bin/python ./setup.py install
-
+cd /home/kin/workspace/OpenSceneFlow
mamba activate opensf
```
@@ -119,7 +122,7 @@ Some tips before running the code:
* If you want to use [wandb](wandb.ai), replace all `entity="kth-rpl",` to your own entity otherwise tensorboard will be used locally.
* Set correct data path by passing the config, e.g. `train_data=/home/kin/data/av2/h5py/demo/train val_data=/home/kin/data/av2/h5py/demo/val`.
-And free yourself from trainning, you can download the pretrained weight from [HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow) and we provided the detail `wget` command in each model section.
+And free yourself from trainning, you can download the pretrained weight from [HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow) and we provided the detail `wget` command in each model section. For optimization-based method, it's train-free so you can directly run with [3. Evaluation](#3-evaluation) (check more in the evaluation section).
```bash
mamba activate opensf
@@ -143,6 +146,8 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckp
Extra pakcges needed for SSF model:
```bash
pip install mmengine-lite torch-scatter
+# torch-scatter might not working, then reinstall by:
+pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bpt20cu118-cp310-cp310-linux_x86_64.whl
```
Train SSF with the leaderboard submit config. [Runtime: Around 6 hours in 8x A100 GPUs.]
@@ -194,9 +199,12 @@ You can view Wandb dashboard for the training and evaluation results or upload r
Since in training, we save all hyper-parameters and model checkpoints, the only thing you need to do is to specify the checkpoint path. Remember to set the data path correctly also.
```bash
-# it will directly prints all metric
+# (feed-forward): load ckpt and run it, it will directly prints all metric
python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=val
+# (optimization-based): it might need take really long time, maybe tmux for run it.
+python eval.py model=nsfp
+
# it will output the av2_submit.zip or av2_submit_v2.zip for you to submit to leaderboard
python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=1
python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=2
@@ -238,7 +246,10 @@ evalai challenge 2210 phase 4396 submit --file av2_submit_v2.zip --large --priva
We provide a script to visualize the results of the model also. You can specify the checkpoint path and the data path to visualize the results. The step is quite similar to evaluation.
```bash
+# (feed-forward): load ckpt
python save.py checkpoint=/home/kin/seflow_best.ckpt dataset_path=/home/kin/data/av2/preprocess_v2/sensor/vis
+# (optimization-based): change another model by passing model name.
+python eval.py model=nsfp dataset_path=/home/kin/data/av2/h5py/demo/val
# The output of above command will be like:
Model: DeFlow, Checkpoint from: /home/kin/model_zoo/v2/seflow_best.ckpt
@@ -252,6 +263,11 @@ python tools/visualization.py --res_name 'seflow_best' --data_dir /home/kin/data
https://github.com/user-attachments/assets/f031d1a2-2d2f-4947-a01f-834ed1c146e6
+For exporting easy comparsion with ground truth and other methods, we also provided multi-visulization open3d window:
+```bash
+python tools/visualization.py --mode mul --res_name "['flow', 'seflow_best']" --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
+```
+
Or another way to interact with [rerun](https://github.com/rerun-io/rerun) but please only vis scene by scene, not all at once.
```bash
diff --git a/assets/README.md b/assets/README.md
index bc954b7..837b01a 100644
--- a/assets/README.md
+++ b/assets/README.md
@@ -49,7 +49,7 @@ Then follow [this stackoverflow answers](https://stackoverflow.com/questions/596
3. Then you can build the docker image:
```bash
- cd OpenSceneFlow && docker build -t zhangkin/OpenSceneFlow .
+ cd OpenSceneFlow && docker build -f Dockerfile -t zhangkin/opensf .
```
## Installation
@@ -98,12 +98,5 @@ python -c "from assets.cuda.chamfer3D import nnChamferDis;print('successfully im
Solved by `export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib`
-## Contribute
-
-If you want to contribute to new model, here are tips you can follow:
-1. Dataloader: we believe all data could be process to `.h5`, we named as different scene and inside a scene, the key of each data is timestamp. Check [dataprocess/README.md](../dataprocess/README.md#process) for more details.
-2. Model: All model files can be found [here: src/models](../src/models). You can view deflow and fastflow3d to know how to implement a new model. Don't forget to add to the `__init__.py` [file to import class](../src/models/__init__.py).
-3. Loss: All loss files can be found [here: src/lossfuncs.py](../src/lossfuncs.py). There are three loss functions already inside the file, you can add a new one following the same pattern.
-4. Training: Once you have implemented the model, you can add the model to the config file [here: conf/model](../conf/model) and train the model using the command `python train.py model=your_model_name`. One more note here may: if your res_dict from model output is different, you may need add one pattern in `def training_step` and `def validation_step`.
-
-All others like eval and vis will be changed according to the model you implemented as you follow the above steps.
\ No newline at end of file
+3. torch_scatter problem: `OSError: /home/kin/mambaforge/envs/opensf-v2/lib/python3.10/site-packages/torch_scatter/_version_cpu.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE`
+ Solved by install the torch-cuda version: `pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bpt20cu118-cp310-cp310-linux_x86_64.whl`
diff --git a/assets/cuda/mmcv/scatter_points.py b/assets/cuda/mmcv/scatter_points.py
index 258ca97..a278cbd 100644
--- a/assets/cuda/mmcv/scatter_points.py
+++ b/assets/cuda/mmcv/scatter_points.py
@@ -8,14 +8,24 @@
# from utils import ext_loader
import importlib
-def load_ext(name, funcs):
- ext = importlib.import_module('mmcv.' + name)
- for fun in funcs:
- assert hasattr(ext, fun), f'{fun} miss in module {name}'
- return ext
+def load_ext(possible_names, funcs):
+ """Try loading module from list of possible names, return first matching."""
+ for name in possible_names:
+ try:
+ ext = importlib.import_module('mmcv' + name)
+ missing = [f for f in funcs if not hasattr(ext, f)]
+ if missing:
+ print(f"Missing functions in 'mmcv{name}': {missing}")
+ continue
+ return ext # success
+ except (ModuleNotFoundError, ImportError) as e:
+ print(f"Failed to import mmcv{name}: {e}")
+ raise ImportError(f"Could not load mmcv extension with functions: {funcs}")
+
+# Usage
ext_module = load_ext(
- '_ext',
+ ['', '._ext'],
['dynamic_point_to_voxel_forward', 'dynamic_point_to_voxel_backward'])
diff --git a/assets/cuda/mmcv/setup.py b/assets/cuda/mmcv/setup.py
index 53ed40b..ba9140a 100644
--- a/assets/cuda/mmcv/setup.py
+++ b/assets/cuda/mmcv/setup.py
@@ -8,7 +8,7 @@
version='1.0.1',
ext_modules=[
CUDAExtension(
- name='mmcv._ext',
+ name='mmcv',
sources=[
"/".join(__file__.split("/")[:-1] + ["scatter_points_cuda.cu"]),
"/".join(__file__.split("/")[:-1] + ["scatter_points.cpp"]),
diff --git a/assets/cuda/mmcv/voxelize.py b/assets/cuda/mmcv/voxelize.py
index 10a6d5c..9574557 100644
--- a/assets/cuda/mmcv/voxelize.py
+++ b/assets/cuda/mmcv/voxelize.py
@@ -8,14 +8,25 @@
# from utils import ext_loader
import importlib
-def load_ext(name, funcs):
- ext = importlib.import_module('mmcv.' + name)
- for fun in funcs:
- assert hasattr(ext, fun), f'{fun} miss in module {name}'
- return ext
+def load_ext(possible_names, funcs):
+ """Try loading module from list of possible names, return first matching."""
+ for name in possible_names:
+ try:
+ ext = importlib.import_module('mmcv' + name)
+ missing = [f for f in funcs if not hasattr(ext, f)]
+ if missing:
+ print(f"Missing functions in 'mmcv{name}': {missing}")
+ continue
+ return ext # success
+ except (ModuleNotFoundError, ImportError) as e:
+ print(f"Failed to import mmcv{name}: {e}")
+ raise ImportError(f"Could not load mmcv extension with functions: {funcs}")
+
+# Usage
ext_module = load_ext(
- '_ext', ['dynamic_voxelize_forward', 'hard_voxelize_forward'])
+ ['', '._ext'],
+ ['dynamic_voxelize_forward', 'hard_voxelize_forward'])
class _Voxelization(Function):
diff --git a/conf/eval.yaml b/conf/eval.yaml
index bc8c216..2c65dbe 100644
--- a/conf/eval.yaml
+++ b/conf/eval.yaml
@@ -8,6 +8,8 @@ leaderboard_version: 1 # [1, 2]
supervised_flag: True # [True, False], whether you use any label from the dataset
# no need to change
+save_res_path:
+num_frames: 2
slurm_id: 00000
output: ${model.name}-${slurm_id}
gpus: 1
diff --git a/conf/model/fastnsf.yaml b/conf/model/fastnsf.yaml
new file mode 100644
index 0000000..fd837bc
--- /dev/null
+++ b/conf/model/fastnsf.yaml
@@ -0,0 +1,17 @@
+name: fastnsf
+
+target:
+ _target_: src.models.FastNSF
+ filter_size: 128
+ act_fn: relu
+ layer_size: 8
+ itr_num: 5000
+ lr: 8e-3
+ min_delta: 5e-5
+ grid_factor: 10
+ init_weight: True
+ early_patience: 10 # default is 10, set 30 to get better results
+ point_cloud_range: [-51.2, -51.2, -3, 51.2, 51.2, 3]
+
+val_monitor: val/Dynamic/Mean
+iter_only: True
\ No newline at end of file
diff --git a/conf/model/nsfp.yaml b/conf/model/nsfp.yaml
new file mode 100644
index 0000000..6163b6b
--- /dev/null
+++ b/conf/model/nsfp.yaml
@@ -0,0 +1,15 @@
+name: nsfp
+
+target:
+ _target_: src.models.NSFP
+ filter_size: 128
+ act_fn: relu
+ layer_size: 8
+ itr_num: 5000
+ lr: 8e-3
+ min_delta: 5e-5
+ early_patience: 30 # default is 30, set 100 to get better results (zeroflow setting.)
+ point_cloud_range: [-51.2, -51.2, -3, 51.2, 51.2, 3]
+
+val_monitor: val/Dynamic/Mean
+iter_only: True
\ No newline at end of file
diff --git a/conf/save.yaml b/conf/save.yaml
index de1f5f5..0332482 100644
--- a/conf/save.yaml
+++ b/conf/save.yaml
@@ -4,6 +4,7 @@ res_name: # if None will directly be the `model_name.ckpt` in checkpoint path
# no need to change
+num_frames: 2
defaults:
- hydra: default
- model: deflow
diff --git a/dataprocess/README.md b/dataprocess/README.md
index d83fbcf..7c57411 100644
--- a/dataprocess/README.md
+++ b/dataprocess/README.md
@@ -11,6 +11,8 @@ We've updated the process dataset for:
- [x] Argoverse 2.0: check [here](#argoverse-20). The process script Involved from [DeFlow](https://github.com/KTH-RPL/DeFlow).
- [x] Waymo: check [here](#waymo-dataset). The process script was involved from [SeFlow](https://github.com/KTH-RPL/SeFlow).
- [ ] nuScenes: done coding, public after review. Will be involved later by another paper.
+- [ ] TruckScene: done coding, public after review. Will be involved later by another paper.
+- [ ] ZOD (w/o gt): done coding, public after review. Will be involved later by another paper.
If you want to **use all datasets above**, there is a **specific environment** in [envsftool.yaml](../envsftool.yaml) to install all the necessary packages. As Waymo package have different configuration and conflict with the main environment. Setup through the following command:
diff --git a/dataprocess/extract_nus.py b/dataprocess/extract_nus.py
index e69de29..6736e4c 100644
--- a/dataprocess/extract_nus.py
+++ b/dataprocess/extract_nus.py
@@ -0,0 +1,38 @@
+"""
+#
+# Created: 2024-02-24 10:48
+# Copyright (C) 2024-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen ZHANG (https://kin-zhang.github.io/), Ajinkya Khoche, Peizheng Li
+#
+# Description: Preprocess Data, save as h5df format for faster loading
+# This one is for nuScenes dataset
+#
+"""
+
+import os
+os.environ["OMP_NUM_THREADS"] = "1"
+
+import multiprocessing
+from pathlib import Path
+from multiprocessing import Pool, current_process
+from typing import Optional, Tuple, Dict, Union, Final
+from tqdm import tqdm
+import numpy as np
+import fire, time, h5py
+
+from nuscenes.nuscenes import NuScenes
+from nuscenes.utils import splits
+from nuscenes.utils.geometry_utils import transform_matrix
+from nuscenes.utils.data_classes import Box
+from nuscenes.utils.geometry_utils import points_in_box
+from pyquaternion import Quaternion
+
+import os, sys
+PARENT_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..'))
+sys.path.append(PARENT_DIR)
+from src.utils.mics import create_reading_index
+from src.utils.av2_eval import CATEGORY_TO_INDEX, NusNamMap
+from linefit import ground_seg
+
+BOUNDING_BOX_EXPANSION: Final = 0.2
+GROUNDSEG_config = f"{PARENT_DIR}/conf/others/nuscenes.toml"
\ No newline at end of file
diff --git a/environment.yaml b/environment.yaml
index b7acc5b..d520844 100644
--- a/environment.yaml
+++ b/environment.yaml
@@ -3,12 +3,13 @@ channels:
- conda-forge
- pytorch
dependencies:
- - python=3.8
+ - python=3.10
- pytorch::pytorch=2.0.0
- pytorch::torchvision
- - pytorch::pytorch-cuda=11.7
- - nvidia/label/cuda-11.7.0::cuda
- - lightning==2.0.1
+ - pytorch::pytorch-cuda=11.8
+ - conda-forge::lightning=2.0.1
+ - nvidia/label/cuda-11.8.0::cuda
+ - nvidia/label/cuda-11.8.0::cuda-toolkit
- mkl==2024.0.0
- tensorboard
- numba
@@ -23,7 +24,7 @@ dependencies:
- hydra-core
- fire
- tabulate
- - scikit-learn==1.3.2
+ - scikit-learn
- hdbscan
- setuptools==69.5.1
- gxx_linux-64==11.4.0
@@ -33,19 +34,18 @@ dependencies:
- assets/cuda/chamfer3D
- assets/cuda/mmcv
- open3d==0.18.0
+ - av2==0.3.1
+ - spconv-cu118==2.3.6
+ - numpy==1.26.4
+ # Qingwen's public pkg:
- dztimer
- - av2==0.2.1
- - dufomap==1.0.0
- - spconv-cu117==2.3.6
- # - mmengine-lite
- # - torch-scatter==2.1.2
-
+ - dufomap==1.1.0
+ - linefit==1.1.0
+
# Reason about the version fixed:
# setuptools==68.5.1: https://github.com/aws-neuron/aws-neuron-sdk/issues/893
# mkl==2024.0.0: https://github.com/pytorch/pytorch/issues/123097#issue-2218541307
# av2==0.2.1: in case other version deleted some functions.
# lightning==2.0.1: https://stackoverflow.com/questions/76647518/how-to-fix-error-cannot-import-name-modelmetaclass-from-pydantic-main
# open3d==0.18.0: because 0.17.0 have bug on set the view json file
-# dufomap==1.0.0: in case later updating may not compatible with the code.
-# spconv-cu117==2.3.6: avoid error: KeyError: ((16, 8, 8), float, float)
-# torch-scatter==2.1.2: in case later updating may not compatible with the code.
+# dufomap==1.1.0: in case later updating may not compatible with the code.
diff --git a/envsftool.yaml b/envsftool.yaml
index 57b5dd7..8ac1925 100644
--- a/envsftool.yaml
+++ b/envsftool.yaml
@@ -21,11 +21,14 @@ dependencies:
- nuscenes-devkit
- av2==0.2.1
- waymo-open-dataset-tf-2.11.0==1.5.0
- - linefit
+ - zod==0.5.0
+ - truckscenes-devkit
+ # Qingwen's public pkg:
- dztimer
- - dufomap==1.0.0
+ - dufomap==1.1.0
+ - linefit==1.1.0
# Reason about the version fixed:
# numpy==1.22: package conflicts, need numpy higher or same 1.22
-# dufomap==1.0.0: in case later updating may not compatible with the code.
+# dufomap==1.1.0: in case later updating may not compatible with the code.
# Pleaase check dataprocess/README.md for more details.
\ No newline at end of file
diff --git a/eval.py b/eval.py
index c2449a5..bfe7bdd 100644
--- a/eval.py
+++ b/eval.py
@@ -33,6 +33,12 @@ def precheck_cfg_valid(cfg):
def main(cfg):
pl.seed_everything(cfg.seed, workers=True)
output_dir = HydraConfig.get().runtime.output_dir
+
+ if 'iter_only' in cfg.model and cfg.model.iter_only:
+ from src.runner import launch_runner
+ print(f"---LOG[eval]: Run optmization-based method: {cfg.model.name}")
+ launch_runner(cfg, cfg.av2_mode)
+ return
if not os.path.exists(cfg.checkpoint):
print(f"Checkpoint {cfg.checkpoint} does not exist. Need checkpoints for evaluation.")
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..cbce091
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,32 @@
+# A minimal env for training env only.
+# --index-url https://download.pytorch.org/whl/cu118
+# torch==2.0.0
+# torchvision==0.15.1
+# torchaudio==2.0.1
+
+lightning
+tensorboard
+wandb
+tqdm
+fire
+omegaconf
+hydra-core
+tabulate
+pathtools
+numpy==1.26.0
+scipy
+numba
+pandas
+scikit-learn
+h5py
+hdbscan
+av2==0.3.4
+spconv-cu118==2.3.8
+flash-attn==2.6.3
+mmengine-lite
+torch-scatter==2.1.2
+
+# Qingwen's public pkg
+dufomap==1.1.0
+linefit==1.1.0
+dztimer
diff --git a/save.py b/save.py
index 00bf259..b9195e6 100644
--- a/save.py
+++ b/save.py
@@ -27,6 +27,13 @@ def main(cfg):
pl.seed_everything(cfg.seed, workers=True)
output_dir = HydraConfig.get().runtime.output_dir
+ if 'iter_only' in cfg.model and cfg.model.iter_only:
+ from src.runner import launch_runner
+ print(f"---LOG[eval]: Run optmization-based method: {cfg.model.name}")
+ cfg.res_name = cfg.model.name if (cfg.res_name is None) else cfg.res_name
+ launch_runner(cfg, 'save')
+ return
+
if not os.path.exists(cfg.checkpoint):
print(f"Checkpoint {cfg.checkpoint} does not exist. Need checkpoints for evaluation.")
sys.exit(1)
diff --git a/src/dataset.py b/src/dataset.py
index fb88d7d..acb514a 100644
--- a/src/dataset.py
+++ b/src/dataset.py
@@ -19,6 +19,9 @@
BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' ))
sys.path.append(BASE_DIR)
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
def collate_fn_pad(batch):
num_frames = 2
@@ -83,6 +86,8 @@ def __init__(self, directory, n_frames=2, dufo=False, eval = False, leaderboard_
super(HDF5Dataset, self).__init__()
self.directory = directory
+ if (torch.distributed.is_initialized() and torch.distributed.get_rank() == 0) or not torch.distributed.is_initialized():
+ print(f"----[Debug] Loading data with num_frames={n_frames}, eval={eval}, leaderboard_version={leaderboard_version}")
with open(os.path.join(self.directory, 'index_total.pkl'), 'rb') as f:
self.data_index = pickle.load(f)
diff --git a/src/models/__init__.py b/src/models/__init__.py
index 503a046..4735f1c 100644
--- a/src/models/__init__.py
+++ b/src/models/__init__.py
@@ -12,6 +12,7 @@
from .deflow import DeFlow
from .fastflow3d import FastFlow3D
+from .nsfp import NSFP
# following need install extra package:
# * pip install spconv-cu117
@@ -27,4 +28,11 @@
from .ssf import SSF
except ImportError as e:
print("\033[93m--- WARNING [model]: Model with torch scatter is not imported, as it requires some lib which is not installed.")
- print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.")
\ No newline at end of file
+ print(f"Detail error message\033[0m: {e}. Just ignore this warning if code runs without these models.")
+
+# following need install extra package:
+# * pip install FastGeodis --no-build-isolation --no-cache-dir
+try:
+ from .fastnsf import FastNSF
+except ImportError:
+ print("--- WARNING [model]: FastNSF is not imported, as it requires FastGeodis lib which is not installed.")
diff --git a/src/models/basic/nsfp_module.py b/src/models/basic/nsfp_module.py
index 8d293f8..bf88fe7 100644
--- a/src/models/basic/nsfp_module.py
+++ b/src/models/basic/nsfp_module.py
@@ -26,6 +26,17 @@ def __init__(self, dim_x=3, filter_size=128, act_fn='relu', layer_size=8):
self.nn_layers.append(torch.nn.Linear(filter_size, dim_x))
else:
self.nn_layers.append(torch.nn.Sequential(torch.nn.Linear(dim_x, dim_x)))
+
+ def reset(self):
+ for layer in self.nn_layers.children():
+ if hasattr(layer, 'reset_parameters'):
+ layer.reset_parameters()
+
+ def init_weights(self):
+ for m in self.nn_layers:
+ if isinstance(m, torch.nn.Linear):
+ torch.nn.init.xavier_uniform_(m.weight)
+ m.bias.data.fill_(0.0)
def forward(self, x):
""" points -> features
@@ -86,210 +97,3 @@ def _init_is_better(self, mode, min_delta, percentage):
self.is_better = lambda a, best: a > best + (
best * min_delta / 100)
-
-
-import torch.nn.functional as F
-from pytorch3d.ops.knn import knn_gather, knn_points
-from pytorch3d.structures.pointclouds import Pointclouds
-from typing import Union
-
-
-def _validate_chamfer_reduction_inputs(batch_reduction: Union[str, None],
- point_reduction: str):
- """Check the requested reductions are valid.
- Args:
- batch_reduction: Reduction operation to apply for the loss across the
- batch, can be one of ["mean", "sum"] or None.
- point_reduction: Reduction operation to apply for the loss across the
- points, can be one of ["mean", "sum"].
- """
- if batch_reduction is not None and batch_reduction not in ["mean", "sum"]:
- raise ValueError(
- 'batch_reduction must be one of ["mean", "sum"] or None')
- if point_reduction not in ["mean", "sum"]:
- raise ValueError('point_reduction must be one of ["mean", "sum"]')
-
-
-def _handle_pointcloud_input(
- points: Union[torch.Tensor, Pointclouds],
- lengths: Union[torch.Tensor, None],
-):
- """
- If points is an instance of Pointclouds, retrieve the padded points tensor
- along with the number of points per batch and the padded normals.
- Otherwise, return the input points (and normals) with the number of points per cloud
- set to the size of the second dimension of `points`.
- """
- if isinstance(points, Pointclouds):
- X = points.points_padded()
- lengths = points.num_points_per_cloud()
- elif torch.is_tensor(points):
- if points.ndim != 3:
- raise ValueError(
- f"Expected points to be of shape (N, P, D), got {points.shape}"
- )
- X = points
- if lengths is not None and (lengths.ndim != 1
- or lengths.shape[0] != X.shape[0]):
- raise ValueError("Expected lengths to be of shape (N,)")
- if lengths is None:
- lengths = torch.full((X.shape[0], ),
- X.shape[1],
- dtype=torch.int64,
- device=points.device)
- else:
- raise ValueError("The input pointclouds should be either " +
- "Pointclouds objects or torch.Tensor of shape " +
- "(minibatch, num_points, 3).")
- return X, lengths
-
-def my_chamfer_fn(
- x,
- y,
- x_lengths=None,
- y_lengths=None,
- x_normals=None,
- y_normals=None,
- weights=None,
- batch_reduction: Union[str, None] = "mean",
- point_reduction: str = "mean",
- truncate_dist=True,
-):
- """
- Chamfer distance between two pointclouds x and y.
-
- Args:
- x: FloatTensor of shape (N, P1, D) or a Pointclouds object representing
- a batch of point clouds with at most P1 points in each batch element,
- batch size N and feature dimension D.
- y: FloatTensor of shape (N, P2, D) or a Pointclouds object representing
- a batch of point clouds with at most P2 points in each batch element,
- batch size N and feature dimension D.
- x_lengths: Optional LongTensor of shape (N,) giving the number of points in each
- cloud in x.
- y_lengths: Optional LongTensor of shape (N,) giving the number of points in each
- cloud in x.
- x_normals: Optional FloatTensor of shape (N, P1, D).
- y_normals: Optional FloatTensor of shape (N, P2, D).
- weights: Optional FloatTensor of shape (N,) giving weights for
- batch elements for reduction operation.
- batch_reduction: Reduction operation to apply for the loss across the
- batch, can be one of ["mean", "sum"] or None.
- point_reduction: Reduction operation to apply for the loss across the
- points, can be one of ["mean", "sum"].
-
- Returns:
- 2-element tuple containing
-
- - **loss**: Tensor giving the reduced distance between the pointclouds
- in x and the pointclouds in y.
- - **loss_normals**: Tensor giving the reduced cosine distance of normals
- between pointclouds in x and pointclouds in y. Returns None if
- x_normals and y_normals are None.
- """
- _validate_chamfer_reduction_inputs(batch_reduction, point_reduction)
-
- x, x_lengths = _handle_pointcloud_input(x, None)
- y, y_lengths = _handle_pointcloud_input(y, None)
-
- if x_lengths.item() == 0 or y_lengths.item() == 0:
- return 0.0, 0.0
- # assert x_lengths.item() > 0, f"x_lengths is {x_lengths.item()}"
- # assert y_lengths.item() > 0, f"y_lengths is {y_lengths.item()}"
-
- return_normals = x_normals is not None and y_normals is not None
-
- N, P1, D = x.shape
- P2 = y.shape[1]
-
- # Check if inputs are heterogeneous and create a lengths mask.
- is_x_heterogeneous = (x_lengths != P1).any()
- is_y_heterogeneous = (y_lengths != P2).any()
- x_mask = (torch.arange(P1, device=x.device)[None] >= x_lengths[:, None]
- ) # shape [N, P1]
- y_mask = (torch.arange(P2, device=y.device)[None] >= y_lengths[:, None]
- ) # shape [N, P2]
-
- if y.shape[0] != N or y.shape[2] != D:
- raise ValueError("y does not have the correct shape.")
- if weights is not None:
- if weights.size(0) != N:
- raise ValueError("weights must be of shape (N,).")
- if not (weights >= 0).all():
- raise ValueError("weights cannot be negative.")
- if weights.sum() == 0.0:
- weights = weights.view(N, 1)
- if batch_reduction in ["mean", "sum"]:
- return (
- (x.sum((1, 2)) * weights).sum() * 0.0,
- (x.sum((1, 2)) * weights).sum() * 0.0,
- )
- return ((x.sum((1, 2)) * weights) * 0.0, (x.sum(
- (1, 2)) * weights) * 0.0)
-
- cham_norm_x = x.new_zeros(())
- cham_norm_y = x.new_zeros(())
-
- x_nn = knn_points(x, y, lengths1=x_lengths, lengths2=y_lengths, K=1)
- y_nn = knn_points(y, x, lengths1=y_lengths, lengths2=x_lengths, K=1)
-
- cham_x = x_nn.dists[..., 0] # (N, P1)
- cham_y = y_nn.dists[..., 0] # (N, P2)
-
- # NOTE: truncated Chamfer distance.
- if truncate_dist:
- dist_thd = 2
- x_mask[cham_x >= dist_thd] = True
- y_mask[cham_y >= dist_thd] = True
- cham_x[x_mask] = 0.0
- cham_y[y_mask] = 0.0
-
- if is_x_heterogeneous:
- cham_x[x_mask] = 0.0
- if is_y_heterogeneous:
- cham_y[y_mask] = 0.0
-
- if weights is not None:
- cham_x *= weights.view(N, 1)
- cham_y *= weights.view(N, 1)
-
- if return_normals:
- # Gather the normals using the indices and keep only value for k=0
- x_normals_near = knn_gather(y_normals, x_nn.idx, y_lengths)[..., 0, :]
- y_normals_near = knn_gather(x_normals, y_nn.idx, x_lengths)[..., 0, :]
-
- cham_norm_x = 1 - torch.abs(
- F.cosine_similarity(x_normals, x_normals_near, dim=2, eps=1e-6))
- cham_norm_y = 1 - torch.abs(
- F.cosine_similarity(y_normals, y_normals_near, dim=2, eps=1e-6))
-
- if is_x_heterogeneous:
- # pyre-fixme[16]: `int` has no attribute `__setitem__`.
- cham_norm_x[x_mask] = 0.0
- if is_y_heterogeneous:
- cham_norm_y[y_mask] = 0.0
-
- if weights is not None:
- cham_norm_x *= weights.view(N, 1)
- cham_norm_y *= weights.view(N, 1)
-
- # Apply point reduction
- cham_x = cham_x.sum(1) # (N,)
- cham_y = cham_y.sum(1) # (N,)
-
- if point_reduction == "mean":
- cham_x /= x_lengths
- cham_y /= y_lengths
-
- if batch_reduction is not None:
- # batch_reduction == "sum"
- cham_x = cham_x.sum()
- cham_y = cham_y.sum()
- if batch_reduction == "mean":
- cham_x /= N
- cham_y /= N
-
- cham_dist = cham_x + cham_y
- cham_normals = cham_norm_x + cham_norm_y if return_normals else None
-
- return cham_dist, cham_normals
\ No newline at end of file
diff --git a/src/models/fastnsf.py b/src/models/fastnsf.py
new file mode 100644
index 0000000..f69ac0e
--- /dev/null
+++ b/src/models/fastnsf.py
@@ -0,0 +1,223 @@
+
+"""
+This file is from: https://github.com/Lilac-Lee/FastNSF
+with slightly modification to have unified format with all benchmark.
+
+# Created: 2024-07-27 11:40
+# Copyright (C) 2024-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen Zhang (https://kin-zhang.github.io/)
+#
+# This file is part of
+# * SeFlow (https://github.com/KTH-RPL/SeFlow)
+# * HiMo (https://kin-zhang.github.io/HiMo)
+# If you find this repo helpful, please cite the respective publication as
+# listed on the above website.
+#
+# Description: FastNSF to our codebase implementation.
+# one more package need install: pip install FastGeodis==1.0.4 --no-build-isolation
+"""
+
+import dztimer, torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import FastGeodis # extra package
+
+from .basic.nsfp_module import Neural_Prior, EarlyStopping
+from .basic import cal_pose0to1
+
+class DT:
+ def __init__(self, pts, pmin, pmax, grid_factor, device='cuda:0'):
+ self.device = device
+ self.grid_factor = grid_factor
+
+ sample_x = ((pmax[0] - pmin[0]) * grid_factor).ceil().int() + 2
+ sample_y = ((pmax[1] - pmin[1]) * grid_factor).ceil().int() + 2
+ sample_z = ((pmax[2] - pmin[2]) * grid_factor).ceil().int() + 2
+
+ self.Vx = torch.linspace(0, sample_x, sample_x+1, device=self.device)[:-1] / grid_factor + pmin[0]
+ self.Vy = torch.linspace(0, sample_y, sample_y+1, device=self.device)[:-1] / grid_factor + pmin[1]
+ self.Vz = torch.linspace(0, sample_z, sample_z+1, device=self.device)[:-1] / grid_factor + pmin[2]
+
+ # NOTE: build a binary image first, with 0-value occuppied points
+ grid_x, grid_y, grid_z = torch.meshgrid(self.Vx, self.Vy, self.Vz, indexing="ij")
+ self.grid = torch.stack([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1), grid_z.unsqueeze(-1)], -1).float().squeeze()
+ H, W, D, _ = self.grid.size()
+ pts_mask = torch.ones(H, W, D, device=device)
+ self.pts_sample_idx_x = ((pts[:,0:1] - self.Vx[0]) * self.grid_factor).round()
+ self.pts_sample_idx_y = ((pts[:,1:2] - self.Vy[0]) * self.grid_factor).round()
+ self.pts_sample_idx_z = ((pts[:,2:3] - self.Vz[0]) * self.grid_factor).round()
+ pts_mask[self.pts_sample_idx_x.long(), self.pts_sample_idx_y.long(), self.pts_sample_idx_z.long()] = 0.
+
+ iterations = 1
+ image_pts = torch.zeros(H, W, D, device=device).unsqueeze(0).unsqueeze(0)
+ pts_mask = pts_mask.unsqueeze(0).unsqueeze(0)
+ self.D = FastGeodis.generalised_geodesic3d(
+ image_pts, pts_mask, [1./self.grid_factor, 1./self.grid_factor, 1./self.grid_factor], 1e10, 0.0, iterations
+ ).squeeze()
+
+ def torch_bilinear_distance(self, Y):
+ H, W, D = self.D.size()
+ target = self.D[None, None, ...]
+
+ sample_x = ((Y[:,0:1] - self.Vx[0]) * self.grid_factor).clip(0, H-1)
+ sample_y = ((Y[:,1:2] - self.Vy[0]) * self.grid_factor).clip(0, W-1)
+ sample_z = ((Y[:,2:3] - self.Vz[0]) * self.grid_factor).clip(0, D-1)
+
+ sample = torch.cat([sample_x, sample_y, sample_z], -1)
+
+ # NOTE: normalize samples to [-1, 1]
+ sample = 2 * sample
+ sample[...,0] = sample[...,0] / (H-1)
+ sample[...,1] = sample[...,1] / (W-1)
+ sample[...,2] = sample[...,2] / (D-1)
+ sample = sample -1
+
+ sample_ = torch.cat([sample[...,2:3], sample[...,1:2], sample[...,0:1]], -1)
+
+ # NOTE: reshape to match 5D volumetric input
+ dist = F.grid_sample(target, sample_.view(1,-1,1,1,3), mode="bilinear", align_corners=True).view(-1)
+ return dist
+
+class FastNSF(nn.Module):
+ def __init__(self, filter_size=128, act_fn='relu', layer_size=8, grid_factor=10.,\
+ itr_num=5000, lr=8e-3, min_delta=0.00005, early_patience=30,
+ verbose=False, point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], init_weight = True):
+ super().__init__()
+
+ self.filter_size = filter_size
+ self.act_fn = act_fn
+ self.layer_size = layer_size
+
+ self.grid_factor = grid_factor # grid cell size=1/grid_factor.
+
+ self.iteration_num = itr_num
+ self.min_delta = min_delta
+ self.lr = lr
+ self.early_patience = early_patience
+ self.verbose = verbose
+ self.point_cloud_range = point_cloud_range
+ self.timer = dztimer.Timing()
+ self.timer.start("NSFP Model Inference")
+ self.init_weight = init_weight
+ print(f"\n---LOG [model]: FastNSF setup itr_num: {itr_num}, lr: {lr}, early_patience: {early_patience}.")
+
+ def optimize(self, dict2loss):
+ device = dict2loss['pc0'].device
+
+ # NOTE(Qingwen): don't know why, but it must be initialized every optimization time.
+ self.timer[5].start("Network Initialization")
+ net = Neural_Prior(filter_size=self.filter_size, act_fn=self.act_fn, layer_size=self.layer_size)
+ net = net.to(device)
+ if self.init_weight:
+ net.init_weights()
+ net.train()
+ self.timer[5].stop()
+
+ pc0 = dict2loss['pc0']
+ pc1 = dict2loss['pc1']
+
+ pc1_min = torch.min(pc0.squeeze(0), 0)[0]
+ pc2_min = torch.min(pc1.squeeze(0), 0)[0]
+ pc1_max = torch.max(pc0.squeeze(0), 0)[0]
+ pc2_max = torch.max(pc1.squeeze(0), 0)[0]
+
+ xmin_int, ymin_int, zmin_int = torch.floor(torch.where(pc1_minpc2_max, pc1_max, pc2_max)* self.grid_factor+1) / self.grid_factor
+ # print('xmin: {}, xmax: {}, ymin: {}, ymax: {}, zmin: {}, zmax: {}'.format(xmin_int, xmax_int, ymin_int, ymax_int, zmin_int, zmax_int))
+
+ # NOTE: build DT map
+ dt = DT(pc1.clone().squeeze(0).to(device), (xmin_int, ymin_int, zmin_int), (xmax_int, ymax_int, zmax_int), self.grid_factor, device)
+ params = net.parameters()
+ best_forward = {'loss': torch.inf}
+
+ optimizer = torch.optim.Adam(params, lr=self.lr, weight_decay=0)
+ early_stopping = EarlyStopping(patience=self.early_patience, min_delta=self.min_delta)
+
+ for itr_ in range(self.iteration_num):
+ optimizer.zero_grad()
+ self.timer[1].start("Network Time")
+
+ self.timer[1][0].start("Forward")
+ forward_flow = net(pc0)
+ self.timer[1][0].stop()
+ pc0_to_pc1 = pc0 + forward_flow
+ self.timer[1].stop()
+
+ self.timer[2].start("loss")
+ loss = dt.torch_bilinear_distance(pc0_to_pc1.squeeze(0)).mean()
+ self.timer[2].stop()
+
+ if loss <= best_forward['loss']:
+ best_forward['loss'] = loss.item()
+ best_forward['flow'] = pc0_to_pc1 - pc0
+
+ if early_stopping.step(loss) and 'flow' in best_forward: # at least one step
+ break
+
+ self.timer[3].start("Loss Backward")
+ loss.backward()
+ self.timer[3].stop()
+
+ self.timer[4].start("Optimizer Step")
+ optimizer.step()
+ self.timer[4].stop()
+
+ if self.verbose:
+ self.timer.print(random_colors=True, bold=True)
+
+ return best_forward
+
+ def range_limit_(self, pc):
+ """
+ Limit the point cloud to the given range.
+ """
+ mask = (pc[:, 0] >= self.point_cloud_range[0]) & (pc[:, 0] <= self.point_cloud_range[3]) & \
+ (pc[:, 1] >= self.point_cloud_range[1]) & (pc[:, 1] <= self.point_cloud_range[4]) & \
+ (pc[:, 2] >= self.point_cloud_range[2]) & (pc[:, 2] <= self.point_cloud_range[5])
+ return pc[mask], mask
+
+ def forward(self, batch):
+ batch_sizes = len(batch["pose0"])
+
+ pose_flows = []
+ batch_final_flow = []
+ for batch_id in range(batch_sizes):
+ self.timer[0].start("Data Processing")
+ pc0 = batch["pc0"][batch_id]
+ pc1 = batch["pc1"][batch_id]
+ selected_pc0, rm0 = self.range_limit_(pc0)
+ selected_pc1, rm1 = self.range_limit_(pc1)
+ self.timer[0][0].start("pose")
+ if 'ego_motion' in batch:
+ pose_0to1 = batch['ego_motion'][batch_id]
+ else:
+ pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id])
+ self.timer[0][0].stop()
+
+ self.timer[0][1].start("transform")
+ # transform selected_pc0 to pc1
+ transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3]
+ self.timer[0][1].stop()
+ pose_flows.append(transform_pc0 - selected_pc0)
+
+ self.timer[0].stop()
+
+ # since pl in val and test mode will disable_grad.
+ with torch.inference_mode(False):
+ with torch.enable_grad():
+ dict2loss = {
+ 'pc0': transform_pc0.clone().detach(), #.requires_grad_(True),
+ 'pc1': selected_pc1.clone().detach() #.requires_grad_(True)
+ }
+ model_res = self.optimize(dict2loss)
+
+ final_flow = torch.zeros_like(pc0)
+ final_flow[rm0] = model_res["flow"].clone().detach().requires_grad_(False)
+ batch_final_flow.append(final_flow)
+
+ res_dict = {"flow": batch_final_flow,
+ "pose_flow": pose_flows
+ }
+
+ return res_dict
\ No newline at end of file
diff --git a/src/models/flow4d.py b/src/models/flow4d.py
index a9c2fdf..b0f52b5 100644
--- a/src/models/flow4d.py
+++ b/src/models/flow4d.py
@@ -47,7 +47,6 @@ def load_from_checkpoint(self, ckpt_path):
return self.load_state_dict(state_dict=state_dict, strict=False)
def forward(self, batch):
- #t_deflow_start = time.time()
"""
input: using the batch from dataloader, which is a dict
Detail: [pc0, pc1, pose0, pose1]
diff --git a/src/models/nsfp.py b/src/models/nsfp.py
new file mode 100644
index 0000000..8467ae5
--- /dev/null
+++ b/src/models/nsfp.py
@@ -0,0 +1,185 @@
+
+"""
+This file is from: https://github.com/Lilac-Lee/Neural_Scene_Flow_Prior
+with our modification to have unified format with our codebase running.
+
+# Created: 2024-07-27 11:33
+# Copyright (C) 2024-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen Zhang (https://kin-zhang.github.io/)
+#
+# This file is part of
+# * SeFlow (https://github.com/KTH-RPL/SeFlow)
+# * HiMo (https://kin-zhang.github.io/HiMo)
+# If you find this repo helpful, please cite the respective publication as
+# listed on the above website.
+#
+# Description: NSFP to our codebase implementation.
+
+"""
+
+import dztimer, torch, copy
+import torch.nn as nn
+
+from .basic import cal_pose0to1
+from .basic.nsfp_module import Neural_Prior, EarlyStopping
+from assets.cuda.chamfer3D import nnChamferDis
+MyCUDAChamferDis = nnChamferDis()
+
+class NSFP(nn.Module):
+ def __init__(self, filter_size=128, act_fn='relu', layer_size=8, \
+ itr_num=5000, lr=8e-3, min_delta=0.00005, early_patience=30,
+ verbose=False, point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3]):
+ super().__init__()
+
+ self.filter_size = filter_size
+ self.act_fn = act_fn
+ self.layer_size = layer_size
+
+ self.iteration_num = itr_num
+ self.min_delta = min_delta
+ self.lr = lr
+ self.early_patience = early_patience
+ self.verbose = verbose
+ self.point_cloud_range = point_cloud_range
+ self.timer = dztimer.Timing()
+ self.timer.start("NSFP Model Inference")
+ print(f"\n--- LOG [model]: NSFP set itr_num: {itr_num}, lr: {lr}, early_patience: {early_patience}.")
+
+ def cal_loss(self, dict2loss, net_inv, pc0_to_pc1):
+
+ self.timer[2].start("Loss")
+
+
+ self.timer[1][1].start("Inverse")
+ inverse_flow = net_inv(pc0_to_pc1)
+ self.timer[1][1].stop()
+ est_pc1_to_pc0 = pc0_to_pc1 - inverse_flow
+
+
+ self.timer[2][0].start("Forward Loss")
+ # forward_loss, _ = my_chamfer_fn(pc0_to_pc1.unsqueeze(0), dict2loss['pc1'].unsqueeze(0))
+ forward_loss = MyCUDAChamferDis.truncated_dis(pc0_to_pc1, dict2loss['pc1'])
+ self.timer[2][0].stop()
+
+ self.timer[2][1].start("Inverse Loss")
+ # inverse_loss, _ = my_chamfer_fn(est_pc1_to_pc0.unsqueeze(0), dict2loss['pc0'].unsqueeze(0))
+ inverse_loss = MyCUDAChamferDis.truncated_dis(est_pc1_to_pc0, dict2loss['pc0'])
+ self.timer[2][1].stop()
+ loss = forward_loss + inverse_loss
+
+ self.timer[2].stop()
+
+ return loss
+
+ def optimize(self, dict2loss):
+ device = dict2loss['pc0'].device
+
+ # NOTE(Qingwen): don't know why, but it must be initialized every optimization time.
+ self.timer[5].start("Network Initialization")
+ net = Neural_Prior(filter_size=self.filter_size, act_fn=self.act_fn, layer_size=self.layer_size)
+ net = net.to(device)
+ net.train()
+ net_inv = copy.deepcopy(net)
+ self.timer[5].stop()
+
+ pc0 = dict2loss['pc0']
+ params = [{
+ 'params': net.parameters(),
+ 'lr': self.lr,
+ 'weight_decay': 0
+ }, {
+ 'params': net_inv.parameters(),
+ 'lr': self.lr,
+ 'weight_decay': 0
+ }]
+ best_forward = {'loss': torch.inf}
+
+ optimizer = torch.optim.Adam(params, lr=self.lr, weight_decay=0)
+ early_stopping = EarlyStopping(patience=self.early_patience, min_delta=self.min_delta)
+
+ for itr_ in range(self.iteration_num):
+ optimizer.zero_grad()
+ self.timer[1].start("Network Time")
+
+ self.timer[1][0].start("Forward")
+ forward_flow = net(pc0)
+ self.timer[1][0].stop()
+ pc0_to_pc1 = pc0 + forward_flow
+ self.timer[1].stop()
+
+ dict2loss['est_flow'] = forward_flow
+ loss = self.cal_loss(dict2loss, net_inv, pc0_to_pc1)
+
+ if loss <= best_forward['loss']:
+ best_forward['loss'] = loss.item()
+ best_forward['flow'] = forward_flow
+
+ if early_stopping.step(loss) and 'flow' in best_forward: # at least one step
+ break
+
+ self.timer[3].start("Loss Backward")
+ loss.backward()
+ self.timer[3].stop()
+
+ self.timer[4].start("Optimizer Step")
+ optimizer.step()
+ self.timer[4].stop()
+
+ if self.verbose:
+ self.timer.print(random_colors=True, bold=True)
+
+ return best_forward
+
+ def range_limit_(self, pc):
+ """
+ Limit the point cloud to the given range.
+ """
+ mask = (pc[:, 0] >= self.point_cloud_range[0]) & (pc[:, 0] <= self.point_cloud_range[3]) & \
+ (pc[:, 1] >= self.point_cloud_range[1]) & (pc[:, 1] <= self.point_cloud_range[4]) & \
+ (pc[:, 2] >= self.point_cloud_range[2]) & (pc[:, 2] <= self.point_cloud_range[5])
+ return pc[mask], mask
+
+ def forward(self, batch):
+ batch_sizes = len(batch["pose0"])
+
+ pose_flows = []
+ batch_final_flow = []
+ for batch_id in range(batch_sizes):
+ self.timer[0].start("Data Processing")
+ pc0 = batch["pc0"][batch_id]
+ pc1 = batch["pc1"][batch_id]
+ selected_pc0, rm0 = self.range_limit_(pc0)
+ selected_pc1, rm1 = self.range_limit_(pc1)
+ self.timer[0][0].start("pose")
+ if 'ego_motion' in batch:
+ pose_0to1 = batch['ego_motion'][batch_id]
+ else:
+ pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id])
+ self.timer[0][0].stop()
+
+ self.timer[0][1].start("transform")
+ # transform selected_pc0 to pc1
+ transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3]
+ self.timer[0][1].stop()
+ pose_flows.append(transform_pc0 - selected_pc0)
+
+ self.timer[0].stop()
+
+ # since pl in val and test mode will disable_grad.
+ with torch.inference_mode(False):
+ with torch.enable_grad():
+ dict2loss = {
+ 'pc0': transform_pc0.clone().detach().requires_grad_(True),
+ 'pc1': selected_pc1.clone().detach().requires_grad_(True)
+ }
+ model_res = self.optimize(dict2loss)
+
+ final_flow = torch.zeros_like(pc0)
+ final_flow[rm0] = model_res["flow"].clone().detach().requires_grad_(False)
+ batch_final_flow.append(final_flow)
+
+ res_dict = {"flow": batch_final_flow,
+ "pose_flow": pose_flows
+ }
+
+ return res_dict
\ No newline at end of file
diff --git a/src/runner.py b/src/runner.py
new file mode 100644
index 0000000..98d5644
--- /dev/null
+++ b/src/runner.py
@@ -0,0 +1,336 @@
+"""
+# Created: 2025-08-09 18:59
+# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology
+# Author: Qingwen Zhang (https://kin-zhang.github.io/), AI Assistant (Google AI Studio)
+#
+# This file is part of
+# * SeFlow (https://github.com/KTH-RPL/SeFlow)
+# * HiMo (https://kin-zhang.github.io/HiMo)
+# If you find this repo helpful, please cite the respective publication as
+# listed on the above website.
+#
+# Description: Runner for optimization-based models (e.g., NSFP, FastNSF) that
+# replaces the need for PyTorch Lightning in evaluation/testing.
+# It handles multi-GPU distribution, execution, and result aggregation.
+#
+"""
+
+import os
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.utils.data import DataLoader
+from torch.utils.data import Sampler
+from collections import defaultdict
+from tqdm import tqdm
+from pathlib import Path
+import h5py
+import numpy as np
+from hydra.utils import instantiate
+import shutil
+
+from .dataset import HDF5Dataset
+from .models.basic import cal_pose0to1
+from .utils.eval_metric import OfficialMetrics, evaluate_leaderboard, evaluate_leaderboard_v2, evaluate_ssf
+from .utils.av2_eval import write_output_file
+from .utils.mics import zip_res
+
+class SceneDistributedSampler(Sampler):
+ """
+ A DistributedSampler that distributes data based on scene IDs, not individual indices.
+ This ensures that all frames from a single scene are processed by the same GPU,
+ preventing I/O conflicts when writing to HDF5 files.
+
+ Args:
+ dataset (Dataset): The dataset to sample from. It must have a `data_index`
+ attribute where each element is a list/tuple like [scene_id, timestamp].
+ num_replicas (int, optional): Number of processes participating in distributed training.
+ Defaults to world_size.
+ rank (int, optional): Rank of the current process. Defaults to current rank.
+ """
+ def __init__(self, dataset, num_replicas=None, rank=None):
+ if num_replicas is None:
+ if not dist.is_available():
+ raise RuntimeError("Requires distributed package to be available")
+ num_replicas = dist.get_world_size()
+ if rank is None:
+ if not dist.is_available():
+ raise RuntimeError("Requires distributed package to be available")
+ rank = dist.get_rank()
+
+ self.dataset = dataset
+ self.num_replicas = num_replicas
+ self.rank = rank
+
+ # 1. Group indices by scene_id
+ scene_to_indices = defaultdict(list)
+ data_index = self.dataset.eval_data_index if self.dataset.eval_index else self.dataset.data_index
+ for idx, (scene_id, _) in enumerate(data_index):
+ scene_to_indices[scene_id].append(idx)
+
+ self.scene_to_indices = scene_to_indices
+ self.scenes = sorted(list(self.scene_to_indices.keys()))
+
+ # 2. Distribute the list of scenes among replicas
+ scenes_for_this_rank = self.scenes[self.rank:len(self.scenes):self.num_replicas]
+
+ # 3. Flatten the indices for the assigned scenes
+ self.indices = []
+ for scene_id in scenes_for_this_rank:
+ self.indices.extend(self.scene_to_indices[scene_id])
+
+ self.num_samples = len(self.indices)
+
+ def __iter__(self):
+ return iter(self.indices)
+
+ def __len__(self):
+ return self.num_samples
+
+class InferenceRunner:
+ def __init__(self, cfg, rank, world_size, mode):
+ try:
+ self.model = instantiate(cfg.model.target)
+ except Exception as e:
+ raise RuntimeError(f"Failed to instantiate model with target '{cfg.model.target}'. Error: {e}") from e
+ self.cfg = cfg
+ self.rank = rank
+ self.world_size = world_size
+ self.device = torch.device(f"cuda:{self.rank}")
+ # 'val' for validation, 'test' for test submission, 'save' for h5 writing
+ self.mode = mode
+
+ self.model.to(self.device)
+ self.metrics = OfficialMetrics() if self.mode in ['val', 'eval'] else None
+ self.save_res_path = cfg.get('save_res_path', None)
+
+ def _setup_dataloader(self):
+ if self.mode in ['val', 'test', 'eval']:
+ dataset_path = self.cfg.dataset_path + f"/{self.cfg.av2_mode}"
+ is_eval_mode = True
+ else: # 'save'
+ dataset_path = self.cfg.dataset_path
+ is_eval_mode = False
+
+ dataset = HDF5Dataset(dataset_path,
+ n_frames=self.cfg.num_frames,
+ eval=is_eval_mode,
+ leaderboard_version=self.cfg.leaderboard_version if 'leaderboard_version' in self.cfg else 1)
+
+ sampler = SceneDistributedSampler(dataset, num_replicas=self.world_size, rank=self.rank)
+
+ return DataLoader(dataset,
+ batch_size=1, # One sample per optimization
+ sampler=sampler,
+ num_workers=self.cfg.get('num_workers', 4),
+ pin_memory=True)
+
+
+ def _process_step(self, batch):
+ # Move all tensors to the appropriate device
+ for key, val in batch.items():
+ if isinstance(val, torch.Tensor):
+ batch[key] = val.to(self.device)
+ elif isinstance(val, list) and len(val) > 0 and isinstance(val[0], torch.Tensor):
+ batch[key] = [item.to(self.device) for item in val]
+
+ batch['origin_pc0'] = batch['pc0'].clone()
+ batch['pc0'] = batch['pc0'][~batch['gm0']].unsqueeze(0)
+ batch['pc1'] = batch['pc1'][~batch['gm1']].unsqueeze(0)
+ self.model.timer[12].start("One Scan")
+ res_dict = self.model(batch)
+ self.model.timer[12].stop()
+
+ # NOTE (Qingwen): Since val and test, we will force set batch_size = 1
+ batch = {key: batch[key][0] for key in batch if len(batch[key])>0}
+ res_dict = {key: res_dict[key][0] for key in res_dict if (res_dict[key]!=None and len(res_dict[key])>0) }
+
+ pc0 = batch['origin_pc0']
+ pose_0to1 = cal_pose0to1(batch["pose0"], batch["pose1"])
+ transform_pc0 = pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3]
+ pose_flow = transform_pc0 - pc0
+
+ final_flow = pose_flow.clone()
+ final_flow[~batch['gm0']] = res_dict['flow'] + pose_flow[~batch['gm0']]
+
+ if self.mode in ['val', 'eval']:
+ eval_mask = batch['eval_mask'].squeeze()
+ gt_flow = batch["flow"]
+ v1_dict = evaluate_leaderboard(final_flow[eval_mask], pose_flow[eval_mask], pc0[eval_mask], \
+ gt_flow[eval_mask], batch['flow_is_valid'][eval_mask], \
+ batch['flow_category_indices'][eval_mask])
+ v2_dict = evaluate_leaderboard_v2(final_flow[eval_mask], pose_flow[eval_mask], pc0[eval_mask],
+ gt_flow[eval_mask], batch['flow_is_valid'][eval_mask],
+ batch['flow_category_indices'][eval_mask])
+ ssf_dict = evaluate_ssf(final_flow[eval_mask], pose_flow[eval_mask], pc0[eval_mask], \
+ gt_flow[eval_mask], batch['flow_is_valid'][eval_mask], batch['flow_category_indices'][eval_mask])
+ self.metrics.step(v1_dict, v2_dict, ssf_dict)
+
+ elif self.mode == 'test':
+ eval_mask = batch['eval_mask'].squeeze()
+ save_pred_flow = final_flow[eval_mask, :3].cpu().detach().numpy()
+ rigid_flow = pose_flow[eval_mask, :3].cpu().detach().numpy()
+ is_dynamic = np.linalg.norm(save_pred_flow - rigid_flow, axis=1, ord=2) >= 0.05
+ sweep_uuid = (batch['scene_id'], batch['timestamp'])
+
+ if self.cfg.leaderboard_version == 2:
+ save_pred_flow = (final_flow - pose_flow).cpu().detach().numpy()
+
+ write_output_file(save_pred_flow, is_dynamic, sweep_uuid, self.save_res_path, self.cfg.leaderboard_version)
+
+ elif self.mode == 'save':
+ key = str(batch['timestamp'])
+ scene_id = batch['scene_id']
+ h5_path = os.path.join(self.cfg.dataset_path, f'{scene_id}.h5')
+ try:
+ with h5py.File(h5_path, 'r+') as f:
+ if self.cfg.res_name in f[key]:
+ del f[key][self.cfg.res_name]
+ f[key].create_dataset(self.cfg.res_name, data=final_flow.cpu().detach().numpy().astype(np.float32))
+ except Exception as e:
+ print(f"Rank {self.rank} failed to write to {h5_path}: {e}")
+
+ def run(self):
+ dataloader = self._setup_dataloader()
+
+
+ iter_bar = tqdm(
+ dataloader,
+ desc=f"[GPU {self.rank}] Initializing...",
+ position=self.rank,
+ leave=False,
+ dynamic_ncols=True
+ )
+
+ for batch in iter_bar:
+ current_scene_id = batch['scene_id'][0]
+
+ iter_bar.set_description(f"[GPU {self.rank}] Processing {current_scene_id}")
+
+ with torch.no_grad():
+ self._process_step(batch)
+
+ iter_bar.close()
+
+ def cleanup(self):
+ dist.destroy_process_group()
+
+
+def _setup_distributed_environment():
+ """
+ Sets up the distributed environment. Prioritizes Slurm variables if available,
+ otherwise assumes torchrun/mp.spawn and initializes based on environment variables.
+ """
+ if 'SLURM_PROCID' in os.environ:
+ # --- Slurm Environment ---
+ rank = int(os.environ['SLURM_PROCID'])
+ world_size = int(os.environ['SLURM_NPROCS'])
+ os.environ['RANK'] = str(rank)
+ os.environ['WORLD_SIZE'] = str(world_size)
+
+ hostnames = os.environ['SLURM_JOB_NODELIST']
+ master_addr = hostnames.split(',')[0].split('(')[0]
+ os.environ['MASTER_ADDR'] = master_addr
+ os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '12355')
+
+ # Let torch.distributed handle the rest, it will use the environment variables
+ dist.init_process_group("nccl")
+ return dist.get_rank(), dist.get_world_size()
+
+
+def _run_process(cfg, mode):
+ rank, world_size = _setup_distributed_environment()
+
+ runner = InferenceRunner(cfg, rank, world_size, mode)
+ runner.run()
+
+ # --- use dist.gather_object to synchronize metrics ---
+ if world_size > 1:
+ gathered_metrics_objects = [None] * world_size if rank == 0 else None
+
+ dist.gather_object(
+ runner.metrics,
+ gathered_metrics_objects if rank == 0 else None,
+ dst=0
+ )
+ else:
+ gathered_metrics_objects = [runner.metrics]
+
+ if rank == 0:
+ if mode in ['val', 'eval']:
+ final_metrics = OfficialMetrics()
+ print(f"\n--- [LOG] Finished processing. Aggregating results from {world_size} GPUs with {len(gathered_metrics_objects)} metrics objects...")
+ for metrics_obj in gathered_metrics_objects:
+ if metrics_obj is None: continue
+
+ for key, val_list in metrics_obj.epe_3way.items():
+ final_metrics.epe_3way[key].extend(val_list)
+
+ for class_idx, class_name in enumerate(metrics_obj.bucketedMatrix.class_names):
+ for range_idx, range_bucket in enumerate(metrics_obj.bucketedMatrix.range_buckets):
+ count = metrics_obj.bucketedMatrix.count_storage_matrix[class_idx, range_idx]
+ if count > 0:
+ avg_epe = metrics_obj.bucketedMatrix.epe_storage_matrix[class_idx, range_idx]
+ avg_range = metrics_obj.bucketedMatrix.range_storage_matrix[class_idx, range_idx]
+ final_metrics.bucketedMatrix.accumulate_value(
+ class_name, range_bucket, avg_epe, avg_range, count
+ )
+ for class_idx, class_name in enumerate(metrics_obj.distanceMatrix.class_names):
+ for range_idx, range_bucket in enumerate(metrics_obj.distanceMatrix.range_buckets):
+ count = metrics_obj.distanceMatrix.count_storage_matrix[class_idx, range_idx]
+ if count > 0:
+ avg_epe = metrics_obj.distanceMatrix.epe_storage_matrix[class_idx, range_idx]
+ avg_range = metrics_obj.distanceMatrix.range_storage_matrix[class_idx, range_idx]
+ final_metrics.distanceMatrix.accumulate_value(
+ class_name, range_bucket, avg_epe, avg_range, count
+ )
+
+ final_metrics.print()
+ else:
+ print(f"\nWe already write the {cfg.res_name} into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:")
+ print(f"python tools/visualization.py --res_name '{cfg.res_name}' --data_dir {cfg.dataset_path}")
+ print(f"Enjoy! ^v^ ------ \n")
+ runner.model.timer.print(random_colors=False, bold=False)
+
+ if world_size > 1:
+ dist.barrier()
+
+ runner.cleanup()
+
+def _spawn_wrapper(rank, world_size, cfg, mode):
+ torch.cuda.set_device(rank)
+
+ os.environ['RANK'] = str(rank)
+ os.environ['WORLD_SIZE'] = str(world_size)
+ os.environ['MASTER_ADDR'] = 'localhost'
+ os.environ['MASTER_PORT'] = cfg.get('master_port', '12355')
+ _run_process(cfg, mode)
+
+def launch_runner(cfg, mode):
+ is_slurm_job = 'SLURM_PROCID' in os.environ
+
+ if not is_slurm_job and not dist.is_initialized():
+ world_size = torch.cuda.device_count()
+ if world_size == 0:
+ raise SystemError("No CUDA devices found.")
+
+ if mode == 'test':
+ cfg.save_res_path = Path(cfg.dataset_path).parent / "results" / cfg.output
+
+ mp.spawn(_spawn_wrapper,
+ args=(world_size, cfg, mode),
+ nprocs=world_size,
+ join=True)
+
+ if mode == 'test':
+ print("\nAll workers finished. Aggregating results into submission file...")
+ final_zip_file = zip_res(
+ cfg.save_res_path,
+ leaderboard_version=cfg.leaderboard_version,
+ is_supervised=False, # all optimization-based now is ssl.
+ output_file=str(cfg.save_res_path) + ".zip"
+ )
+ print(f"--- [LOG] Submission file successfully created at: {final_zip_file}")
+ return
+ else:
+ _run_process(cfg, mode)