MKL version fix in CI (#1820)

Summary:
Fix for "undefined symbol: iJIT_NotifyEvent" build issue,

Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/1820

Differential Revision: D58685326

.circleci/config.in.yml

@@ -64,7 +64,7 @@ jobs:
       CUDA_VERSION: "11.3"
     resource_class: gpu.nvidia.small.multi
     machine:
-      image: ubuntu-2004:202101-01
+      image: linux-cuda-11:default
     steps:
       - checkout
       - <<: *setupcuda
@@ -116,7 +116,7 @@ jobs:
       # so we aren't running the tests.
       - run:
           name: build
-          no_output_timeout: 20m
+          no_output_timeout: 40m
           command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py
       - store_artifacts:
           path: /opt/conda/conda-bld/linux-64
@@ -128,7 +128,7 @@ jobs:
   binary_linux_conda_cuda:
     <<: *binary_common
     machine:
-      image: ubuntu-1604-cuda-10.2:202012-01
+      image: linux-cuda-11:default
     resource_class: gpu.nvidia.small.multi
     steps:
     - checkout
@@ -145,7 +145,7 @@ jobs:
           docker pull $TESTRUN_DOCKER_IMAGE
     - run:
         name: Build and run tests
-        no_output_timeout: 20m
+        no_output_timeout: 40m
         command: |
           set -e
 
@@ -156,24 +156,6 @@ jobs:
 
           docker run --gpus all  --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py
 
-  binary_macos_wheel:
-    <<: *binary_common
-    macos:
-      xcode: "13.4.1"
-    steps:
-      - checkout
-      - run:
-          # Cannot easily deduplicate this as source'ing activate
-          # will set environment variables which we need to propagate
-          # to build_wheel.sh
-          command: |
-            curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
-            sh conda.sh -b
-            source $HOME/miniconda3/bin/activate
-            packaging/build_wheel.sh
-      - store_artifacts:
-          path: dist
-
 workflows:
   version: 2
   build_and_test:
@@ -182,23 +164,8 @@ workflows:
       #     context: DOCKERHUB_TOKEN
       {{workflows()}}
       - binary_linux_conda_cuda:
-          name: testrun_conda_cuda_py38_cu102_pyt190
+          name: testrun_conda_cuda_py310_cu117_pyt201
           context: DOCKERHUB_TOKEN
-          python_version: "3.8"
-          pytorch_version: '1.9.0'
-          cu_version: "cu102"
-      - binary_macos_wheel:
-          cu_version: cpu
-          name: macos_wheel_py3.8_cpu
-          python_version: '3.8'
-          pytorch_version: '1.13.0'
-      - binary_macos_wheel:
-          cu_version: cpu
-          name: macos_wheel_py3.9_cpu
-          python_version: '3.9'
-          pytorch_version: '1.13.0'
-      - binary_macos_wheel:
-          cu_version: cpu
-          name: macos_wheel_py3.10_cpu
-          python_version: '3.10'
-          pytorch_version: '1.13.0'
+          python_version: "3.10"
+          pytorch_version: '2.0.1'
+          cu_version: "cu117"

.circleci/config.yml

@@ -64,7 +64,7 @@ jobs:
       CUDA_VERSION: "11.3"
     resource_class: gpu.nvidia.small.multi
     machine:
-      image: ubuntu-2004:202101-01
+      image: linux-cuda-11:default
     steps:
       - checkout
       - <<: *setupcuda
@@ -116,7 +116,7 @@ jobs:
       # so we aren't running the tests.
       - run:
           name: build
-          no_output_timeout: 20m
+          no_output_timeout: 40m
           command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py
       - store_artifacts:
           path: /opt/conda/conda-bld/linux-64
@@ -128,7 +128,7 @@ jobs:
   binary_linux_conda_cuda:
     <<: *binary_common
     machine:
-      image: ubuntu-1604-cuda-10.2:202012-01
+      image: linux-cuda-11:default
     resource_class: gpu.nvidia.small.multi
     steps:
     - checkout
@@ -145,7 +145,7 @@ jobs:
           docker pull $TESTRUN_DOCKER_IMAGE
     - run:
         name: Build and run tests
-        no_output_timeout: 20m
+        no_output_timeout: 40m
         command: |
           set -e
 
@@ -156,143 +156,12 @@ jobs:
 
           docker run --gpus all  --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py
 
-  binary_macos_wheel:
-    <<: *binary_common
-    macos:
-      xcode: "13.4.1"
-    steps:
-      - checkout
-      - run:
-          # Cannot easily deduplicate this as source'ing activate
-          # will set environment variables which we need to propagate
-          # to build_wheel.sh
-          command: |
-            curl -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
-            sh conda.sh -b
-            source $HOME/miniconda3/bin/activate
-            packaging/build_wheel.sh
-      - store_artifacts:
-          path: dist
-
 workflows:
   version: 2
   build_and_test:
     jobs:
       # - main:
       #     context: DOCKERHUB_TOKEN
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt190
-          python_version: '3.8'
-          pytorch_version: 1.9.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py38_cu111_pyt190
-          python_version: '3.8'
-          pytorch_version: 1.9.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt191
-          python_version: '3.8'
-          pytorch_version: 1.9.1
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py38_cu111_pyt191
-          python_version: '3.8'
-          pytorch_version: 1.9.1
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt1100
-          python_version: '3.8'
-          pytorch_version: 1.10.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py38_cu111_pyt1100
-          python_version: '3.8'
-          pytorch_version: 1.10.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py38_cu113_pyt1100
-          python_version: '3.8'
-          pytorch_version: 1.10.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt1101
-          python_version: '3.8'
-          pytorch_version: 1.10.1
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py38_cu111_pyt1101
-          python_version: '3.8'
-          pytorch_version: 1.10.1
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py38_cu113_pyt1101
-          python_version: '3.8'
-          pytorch_version: 1.10.1
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt1102
-          python_version: '3.8'
-          pytorch_version: 1.10.2
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py38_cu111_pyt1102
-          python_version: '3.8'
-          pytorch_version: 1.10.2
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py38_cu113_pyt1102
-          python_version: '3.8'
-          pytorch_version: 1.10.2
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt1110
-          python_version: '3.8'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py38_cu111_pyt1110
-          python_version: '3.8'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py38_cu113_pyt1110
-          python_version: '3.8'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda115
-          context: DOCKERHUB_TOKEN
-          cu_version: cu115
-          name: linux_conda_py38_cu115_pyt1110
-          python_version: '3.8'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt1120
-          python_version: '3.8'
-          pytorch_version: 1.12.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda113
           context: DOCKERHUB_TOKEN
@@ -307,12 +176,6 @@ workflows:
           name: linux_conda_py38_cu116_pyt1120
           python_version: '3.8'
           pytorch_version: 1.12.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py38_cu102_pyt1121
-          python_version: '3.8'
-          pytorch_version: 1.12.1
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda113
           context: DOCKERHUB_TOKEN
@@ -370,118 +233,103 @@ workflows:
           python_version: '3.8'
           pytorch_version: 2.0.0
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt190
-          python_version: '3.9'
-          pytorch_version: 1.9.0
+          cu_version: cu117
+          name: linux_conda_py38_cu117_pyt201
+          python_version: '3.8'
+          pytorch_version: 2.0.1
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py39_cu111_pyt190
-          python_version: '3.9'
-          pytorch_version: 1.9.0
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt201
+          python_version: '3.8'
+          pytorch_version: 2.0.1
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt191
-          python_version: '3.9'
-          pytorch_version: 1.9.1
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt210
+          python_version: '3.8'
+          pytorch_version: 2.1.0
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py39_cu111_pyt191
-          python_version: '3.9'
-          pytorch_version: 1.9.1
+          cu_version: cu121
+          name: linux_conda_py38_cu121_pyt210
+          python_version: '3.8'
+          pytorch_version: 2.1.0
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt1100
-          python_version: '3.9'
-          pytorch_version: 1.10.0
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt211
+          python_version: '3.8'
+          pytorch_version: 2.1.1
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py39_cu111_pyt1100
-          python_version: '3.9'
-          pytorch_version: 1.10.0
+          cu_version: cu121
+          name: linux_conda_py38_cu121_pyt211
+          python_version: '3.8'
+          pytorch_version: 2.1.1
       - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py39_cu113_pyt1100
-          python_version: '3.9'
-          pytorch_version: 1.10.0
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt212
+          python_version: '3.8'
+          pytorch_version: 2.1.2
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt1101
-          python_version: '3.9'
-          pytorch_version: 1.10.1
+          cu_version: cu121
+          name: linux_conda_py38_cu121_pyt212
+          python_version: '3.8'
+          pytorch_version: 2.1.2
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py39_cu111_pyt1101
-          python_version: '3.9'
-          pytorch_version: 1.10.1
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt220
+          python_version: '3.8'
+          pytorch_version: 2.2.0
       - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py39_cu113_pyt1101
-          python_version: '3.9'
-          pytorch_version: 1.10.1
+          cu_version: cu121
+          name: linux_conda_py38_cu121_pyt220
+          python_version: '3.8'
+          pytorch_version: 2.2.0
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt1102
-          python_version: '3.9'
-          pytorch_version: 1.10.2
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt222
+          python_version: '3.8'
+          pytorch_version: 2.2.2
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py39_cu111_pyt1102
-          python_version: '3.9'
-          pytorch_version: 1.10.2
+          cu_version: cu121
+          name: linux_conda_py38_cu121_pyt222
+          python_version: '3.8'
+          pytorch_version: 2.2.2
       - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py39_cu113_pyt1102
-          python_version: '3.9'
-          pytorch_version: 1.10.2
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt231
+          python_version: '3.8'
+          pytorch_version: 2.3.1
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt1110
-          python_version: '3.9'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py39_cu111_pyt1110
-          python_version: '3.9'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
-          context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py39_cu113_pyt1110
-          python_version: '3.9'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda115
-          context: DOCKERHUB_TOKEN
-          cu_version: cu115
-          name: linux_conda_py39_cu115_pyt1110
-          python_version: '3.9'
-          pytorch_version: 1.11.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt1120
-          python_version: '3.9'
-          pytorch_version: 1.12.0
+          cu_version: cu121
+          name: linux_conda_py38_cu121_pyt231
+          python_version: '3.8'
+          pytorch_version: 2.3.1
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda113
           context: DOCKERHUB_TOKEN
@@ -496,12 +344,6 @@ workflows:
           name: linux_conda_py39_cu116_pyt1120
           python_version: '3.9'
           pytorch_version: 1.12.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py39_cu102_pyt1121
-          python_version: '3.9'
-          pytorch_version: 1.12.1
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda113
           context: DOCKERHUB_TOKEN
@@ -559,37 +401,103 @@ workflows:
           python_version: '3.9'
           pytorch_version: 2.0.0
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py310_cu102_pyt1110
-          python_version: '3.10'
-          pytorch_version: 1.11.0
+          cu_version: cu117
+          name: linux_conda_py39_cu117_pyt201
+          python_version: '3.9'
+          pytorch_version: 2.0.1
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu111
-          name: linux_conda_py310_cu111_pyt1110
-          python_version: '3.10'
-          pytorch_version: 1.11.0
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt201
+          python_version: '3.9'
+          pytorch_version: 2.0.1
       - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda113
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu113
-          name: linux_conda_py310_cu113_pyt1110
-          python_version: '3.10'
-          pytorch_version: 1.11.0
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt210
+          python_version: '3.9'
+          pytorch_version: 2.1.0
       - binary_linux_conda:
-          conda_docker_image: pytorch/conda-builder:cuda115
+          conda_docker_image: pytorch/conda-builder:cuda121
           context: DOCKERHUB_TOKEN
-          cu_version: cu115
-          name: linux_conda_py310_cu115_pyt1110
-          python_version: '3.10'
-          pytorch_version: 1.11.0
+          cu_version: cu121
+          name: linux_conda_py39_cu121_pyt210
+          python_version: '3.9'
+          pytorch_version: 2.1.0
       - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
           context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py310_cu102_pyt1120
-          python_version: '3.10'
-          pytorch_version: 1.12.0
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt211
+          python_version: '3.9'
+          pytorch_version: 2.1.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py39_cu121_pyt211
+          python_version: '3.9'
+          pytorch_version: 2.1.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt212
+          python_version: '3.9'
+          pytorch_version: 2.1.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py39_cu121_pyt212
+          python_version: '3.9'
+          pytorch_version: 2.1.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt220
+          python_version: '3.9'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py39_cu121_pyt220
+          python_version: '3.9'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt222
+          python_version: '3.9'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py39_cu121_pyt222
+          python_version: '3.9'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt231
+          python_version: '3.9'
+          pytorch_version: 2.3.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py39_cu121_pyt231
+          python_version: '3.9'
+          pytorch_version: 2.3.1
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda113
           context: DOCKERHUB_TOKEN
@@ -604,12 +512,6 @@ workflows:
           name: linux_conda_py310_cu116_pyt1120
           python_version: '3.10'
           pytorch_version: 1.12.0
-      - binary_linux_conda:
-          context: DOCKERHUB_TOKEN
-          cu_version: cu102
-          name: linux_conda_py310_cu102_pyt1121
-          python_version: '3.10'
-          pytorch_version: 1.12.1
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda113
           context: DOCKERHUB_TOKEN
@@ -666,24 +568,233 @@ workflows:
           name: linux_conda_py310_cu118_pyt200
           python_version: '3.10'
           pytorch_version: 2.0.0
-      - binary_linux_conda_cuda:
-          name: testrun_conda_cuda_py38_cu102_pyt190
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
           context: DOCKERHUB_TOKEN
-          python_version: "3.8"
-          pytorch_version: '1.9.0'
-          cu_version: "cu102"
-      - binary_macos_wheel:
-          cu_version: cpu
-          name: macos_wheel_py3.8_cpu
-          python_version: '3.8'
-          pytorch_version: '1.13.0'
-      - binary_macos_wheel:
-          cu_version: cpu
-          name: macos_wheel_py3.9_cpu
-          python_version: '3.9'
-          pytorch_version: '1.13.0'
-      - binary_macos_wheel:
-          cu_version: cpu
-          name: macos_wheel_py3.10_cpu
+          cu_version: cu117
+          name: linux_conda_py310_cu117_pyt201
           python_version: '3.10'
-          pytorch_version: '1.13.0'
+          pytorch_version: 2.0.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt201
+          python_version: '3.10'
+          pytorch_version: 2.0.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt210
+          python_version: '3.10'
+          pytorch_version: 2.1.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py310_cu121_pyt210
+          python_version: '3.10'
+          pytorch_version: 2.1.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt211
+          python_version: '3.10'
+          pytorch_version: 2.1.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py310_cu121_pyt211
+          python_version: '3.10'
+          pytorch_version: 2.1.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt212
+          python_version: '3.10'
+          pytorch_version: 2.1.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py310_cu121_pyt212
+          python_version: '3.10'
+          pytorch_version: 2.1.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt220
+          python_version: '3.10'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py310_cu121_pyt220
+          python_version: '3.10'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt222
+          python_version: '3.10'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py310_cu121_pyt222
+          python_version: '3.10'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt231
+          python_version: '3.10'
+          pytorch_version: 2.3.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py310_cu121_pyt231
+          python_version: '3.10'
+          pytorch_version: 2.3.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py311_cu118_pyt210
+          python_version: '3.11'
+          pytorch_version: 2.1.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py311_cu121_pyt210
+          python_version: '3.11'
+          pytorch_version: 2.1.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py311_cu118_pyt211
+          python_version: '3.11'
+          pytorch_version: 2.1.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py311_cu121_pyt211
+          python_version: '3.11'
+          pytorch_version: 2.1.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py311_cu118_pyt212
+          python_version: '3.11'
+          pytorch_version: 2.1.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py311_cu121_pyt212
+          python_version: '3.11'
+          pytorch_version: 2.1.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py311_cu118_pyt220
+          python_version: '3.11'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py311_cu121_pyt220
+          python_version: '3.11'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py311_cu118_pyt222
+          python_version: '3.11'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py311_cu121_pyt222
+          python_version: '3.11'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py311_cu118_pyt231
+          python_version: '3.11'
+          pytorch_version: 2.3.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py311_cu121_pyt231
+          python_version: '3.11'
+          pytorch_version: 2.3.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py312_cu118_pyt220
+          python_version: '3.12'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py312_cu121_pyt220
+          python_version: '3.12'
+          pytorch_version: 2.2.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py312_cu118_pyt222
+          python_version: '3.12'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py312_cu121_pyt222
+          python_version: '3.12'
+          pytorch_version: 2.2.2
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py312_cu118_pyt231
+          python_version: '3.12'
+          pytorch_version: 2.3.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda121
+          context: DOCKERHUB_TOKEN
+          cu_version: cu121
+          name: linux_conda_py312_cu121_pyt231
+          python_version: '3.12'
+          pytorch_version: 2.3.1
+      - binary_linux_conda_cuda:
+          name: testrun_conda_cuda_py310_cu117_pyt201
+          context: DOCKERHUB_TOKEN
+          python_version: "3.10"
+          pytorch_version: '2.0.1'
+          cu_version: "cu117"

.circleci/regenerate.py

@@ -18,25 +18,23 @@ from packaging import version
 
 # The CUDA versions which have pytorch conda packages available for linux for each
 # version of pytorch.
-# Pytorch 1.4 also supports cuda 10.0 but we no longer build for cuda 10.0 at all.
 CONDA_CUDA_VERSIONS = {
-    "1.9.0": ["cu102", "cu111"],
-    "1.9.1": ["cu102", "cu111"],
-    "1.10.0": ["cu102", "cu111", "cu113"],
-    "1.10.1": ["cu102", "cu111", "cu113"],
-    "1.10.2": ["cu102", "cu111", "cu113"],
-    "1.11.0": ["cu102", "cu111", "cu113", "cu115"],
-    "1.12.0": ["cu102", "cu113", "cu116"],
-    "1.12.1": ["cu102", "cu113", "cu116"],
+    "1.12.0": ["cu113", "cu116"],
+    "1.12.1": ["cu113", "cu116"],
     "1.13.0": ["cu116", "cu117"],
     "1.13.1": ["cu116", "cu117"],
     "2.0.0": ["cu117", "cu118"],
+    "2.0.1": ["cu117", "cu118"],
+    "2.1.0": ["cu118", "cu121"],
+    "2.1.1": ["cu118", "cu121"],
+    "2.1.2": ["cu118", "cu121"],
+    "2.2.0": ["cu118", "cu121"],
+    "2.2.2": ["cu118", "cu121"],
+    "2.3.1": ["cu118", "cu121"],
 }
 
 
 def conda_docker_image_for_cuda(cuda_version):
-    if cuda_version in ("cu101", "cu102", "cu111"):
-        return None
     if len(cuda_version) != 5:
         raise ValueError("Unknown cuda version")
     return "pytorch/conda-builder:cuda" + cuda_version[2:]
@@ -51,12 +49,24 @@ def pytorch_versions_for_python(python_version):
             for i in CONDA_CUDA_VERSIONS
             if version.Version(i) >= version.Version("1.11.0")
         ]
+    if python_version == "3.11":
+        return [
+            i
+            for i in CONDA_CUDA_VERSIONS
+            if version.Version(i) >= version.Version("2.1.0")
+        ]
+    if python_version == "3.12":
+        return [
+            i
+            for i in CONDA_CUDA_VERSIONS
+            if version.Version(i) >= version.Version("2.2.0")
+        ]
 
 
 def workflows(prefix="", filter_branch=None, upload=False, indentation=6):
     w = []
     for btype in ["conda"]:
-        for python_version in ["3.8", "3.9", "3.10"]:
+        for python_version in ["3.8", "3.9", "3.10", "3.11", "3.12"]:
             for pytorch_version in pytorch_versions_for_python(python_version):
                 for cu_version in CONDA_CUDA_VERSIONS[pytorch_version]:
                     w += workflow_pair(

.flake8

@@ -1,5 +1,8 @@
 [flake8]
-ignore = E203, E266, E501, W503, E221
+# B028 No explicit stacklevel argument found.
+# B907 'foo' is manually surrounded by quotes, consider using the `!r` conversion flag.
+# B905 `zip()` without an explicit `strict=` parameter.
+ignore = E203, E266, E501, W503, E221, B028, B905, B907
 max-line-length = 88
 max-complexity = 18
 select = B,C,E,F,W,T4,B9

INSTALL.md

@@ -9,7 +9,7 @@ The core library is written in PyTorch. Several components have underlying imple
 
 - Linux or macOS or Windows
 - Python 3.8, 3.9 or 3.10
-- PyTorch 1.9.0, 1.9.1, 1.10.0, 1.10.1, 1.10.2, 1.11.0, 1.12.0, 1.12.1, 1.13.0 or 2.0.0.
+- PyTorch 1.12.0, 1.12.1, 1.13.0, 2.0.0, 2.0.1, 2.1.0, 2.1.1, 2.1.2, 2.2.0, 2.2.1, 2.2.2, 2.3.0 or 2.3.1.
 - torchvision that matches the PyTorch installation. You can install them together as explained at pytorch.org to make sure of this.
 - gcc & g++ ≥ 4.9
 - [fvcore](https://github.com/facebookresearch/fvcore)
@@ -77,13 +77,8 @@ Or, to install a nightly (non-official, alpha) build:
 # Anaconda Cloud
 conda install pytorch3d -c pytorch3d-nightly
 ```
-### 2. Install from PyPI, on Mac only.
-This works with pytorch 1.13.0 only. The build is CPU only.
-```
-pip install pytorch3d
-```
 
-### 3. Install wheels for Linux
+### 2. Install wheels for Linux
 We have prebuilt wheels with CUDA for Linux for PyTorch 1.11.0, for each of the supported CUDA versions,
 for Python 3.8 and 3.9. This is for ease of use on Google Colab.
 These are installed in a special way.

README.md

@@ -146,6 +146,12 @@ If you are using the pulsar backend for sphere-rendering (the `PulsarPointRender
 
 Please see below for a timeline of the codebase updates in reverse chronological order. We are sharing updates on the releases as well as research projects which are built with PyTorch3D. The changelogs for the releases are available under [`Releases`](https://github.com/facebookresearch/pytorch3d/releases),  and the builds can be installed using `conda` as per the instructions in [INSTALL.md](INSTALL.md).
 
+**[Oct 31st 2023]:**   PyTorch3D [v0.7.5](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.5) released.
+
+**[May 10th 2023]:**   PyTorch3D [v0.7.4](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.4) released.
+
+**[Apr 5th 2023]:**   PyTorch3D [v0.7.3](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.3) released.
+
 **[Dec 19th 2022]:**   PyTorch3D [v0.7.2](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.2) released.
 
 **[Oct 23rd 2022]:**   PyTorch3D [v0.7.1](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.1) released.

docs/.readthedocs.yaml

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Set the version of Python and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/conf.py
+
+# We recommend specifying your dependencies to enable reproducible builds:
+# https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
+python:
+  install:
+  - requirements: docs/requirements.txt

docs/README.md

@@ -3,7 +3,7 @@
 ### Install dependencies
 
 ```
-pip install -U recommonmark mock sphinx sphinx_rtd_theme sphinx_markdown_tables
+pip install -U recommonmark sphinx sphinx_rtd_theme sphinx_markdown_tables
 ```
 
 ### Add symlink to the root README.md

docs/conf.py

@@ -20,7 +20,8 @@
 import os
 import sys
 
-import mock
+import unittest.mock as mock
+
 from recommonmark.parser import CommonMarkParser
 from recommonmark.states import DummyStateMachine
 from sphinx.builders.html import StandaloneHTMLBuilder

docs/notes/cameras.md

@@ -85,7 +85,7 @@ cameras_ndc = PerspectiveCameras(focal_length=fcl_ndc, principal_point=prp_ndc)
 # Screen space camera
 image_size = ((128, 256),)    # (h, w)
 fcl_screen = (76.8,)          # fcl_ndc * min(image_size) / 2
-prp_screen = ((115.2, 48), )  # w / 2 - px_ndc * min(image_size) / 2, h / 2 - py_ndc * min(image_size) / 2
+prp_screen = ((115.2, 32), )  # w / 2 - px_ndc * min(image_size) / 2, h / 2 - py_ndc * min(image_size) / 2
 cameras_screen = PerspectiveCameras(focal_length=fcl_screen, principal_point=prp_screen, in_ndc=False, image_size=image_size)
 ```
 

docs/requirements.txt

@@ -1,12 +1,11 @@
 docutils>=0.14
 Sphinx>=1.7
-recommonmark==0.4.0
+recommonmark
 sphinx_rtd_theme
 sphinx_markdown_tables
-mock
 numpy
 iopath
 fvcore
-https://download.pytorch.org/whl/cpu/torchvision-0.8.2%2Bcpu-cp37-cp37m-linux_x86_64.whl
-https://download.pytorch.org/whl/cpu/torch-1.7.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torchvision-0.15.2%2Bcpu-cp311-cp311-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torch-2.0.1%2Bcpu-cp311-cp311-linux_x86_64.whl
 omegaconf

docs/tutorials/bundle_adjustment.ipynb

@@ -83,25 +83,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/camera_position_optimization_with_differentiable_rendering.ipynb

@@ -70,25 +70,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/dataloaders_ShapeNetCore_R2N2.ipynb

@@ -45,25 +45,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {
@@ -405,7 +411,7 @@
    "outputs": [],
    "source": [
     "random_model_images = shapenet_dataset.render(\n",
-    "    sample_nums=[3],\n",
+    "    sample_nums=[5],\n",
     "    device=device,\n",
     "    cameras=cameras,\n",
     "    raster_settings=raster_settings,\n",

docs/tutorials/deform_source_mesh_to_target_mesh.ipynb

@@ -84,25 +84,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {
@@ -262,7 +268,7 @@
     "    points = sample_points_from_meshes(mesh, 5000)\n",
     "    x, y, z = points.clone().detach().cpu().squeeze().unbind(1)    \n",
     "    fig = plt.figure(figsize=(5, 5))\n",
-    "    ax = Axes3D(fig)\n",
+    "    ax = fig.add_subplot(111, projection='3d')\n",
     "    ax.scatter3D(x, z, -y)\n",
     "    ax.set_xlabel('x')\n",
     "    ax.set_ylabel('z')\n",

docs/tutorials/fit_simple_neural_radiance_field.ipynb

@@ -50,25 +50,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/fit_textured_mesh.ipynb

@@ -62,25 +62,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/fit_textured_volume.ipynb

@@ -41,25 +41,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/implicitron_config_system.ipynb

@@ -72,25 +72,31 @@
         "import os\n",
         "import sys\n",
         "import torch\n",
+        "import subprocess\n",
         "need_pytorch3d=False\n",
         "try:\n",
         "    import pytorch3d\n",
         "except ModuleNotFoundError:\n",
         "    need_pytorch3d=True\n",
         "if need_pytorch3d:\n",
-        "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-        "        # We try to install PyTorch3D via a released wheel.\n",
-        "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-        "        version_str=\"\".join([\n",
-        "            f\"py3{sys.version_info.minor}_cu\",\n",
-        "            torch.version.cuda.replace(\".\",\"\"),\n",
-        "            f\"_pyt{pyt_version_str}\"\n",
-        "        ])\n",
-        "        !pip install fvcore iopath\n",
+        "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+        "    version_str=\"\".join([\n",
+        "        f\"py3{sys.version_info.minor}_cu\",\n",
+        "        torch.version.cuda.replace(\".\",\"\"),\n",
+        "        f\"_pyt{pyt_version_str}\"\n",
+        "    ])\n",
+        "    !pip install fvcore iopath\n",
+        "    if sys.platform.startswith(\"linux\"):\n",
+        "        print(\"Trying to install wheel for PyTorch3D\")\n",
         "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-        "    else:\n",
-        "        # We try to install PyTorch3D from source.\n",
-        "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+        "        pip_list = !pip freeze\n",
+        "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+        "    if need_pytorch3d:\n",
+        "        print(f\"failed to find/install wheel for {version_str}\")\n",
+        "if need_pytorch3d:\n",
+        "    print(\"Installing PyTorch3D from source\")\n",
+        "    !pip install ninja\n",
+        "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
       ]
     },
     {

docs/tutorials/implicitron_volumes.ipynb

@@ -66,25 +66,31 @@
         "import os\n",
         "import sys\n",
         "import torch\n",
+        "import subprocess\n",
         "need_pytorch3d=False\n",
         "try:\n",
         "    import pytorch3d\n",
         "except ModuleNotFoundError:\n",
         "    need_pytorch3d=True\n",
         "if need_pytorch3d:\n",
-        "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-        "        # We try to install PyTorch3D via a released wheel.\n",
-        "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-        "        version_str=\"\".join([\n",
-        "            f\"py3{sys.version_info.minor}_cu\",\n",
-        "            torch.version.cuda.replace(\".\",\"\"),\n",
-        "            f\"_pyt{pyt_version_str}\"\n",
-        "        ])\n",
-        "        !pip install fvcore iopath\n",
+        "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+        "    version_str=\"\".join([\n",
+        "        f\"py3{sys.version_info.minor}_cu\",\n",
+        "        torch.version.cuda.replace(\".\",\"\"),\n",
+        "        f\"_pyt{pyt_version_str}\"\n",
+        "    ])\n",
+        "    !pip install fvcore iopath\n",
+        "    if sys.platform.startswith(\"linux\"):\n",
+        "        print(\"Trying to install wheel for PyTorch3D\")\n",
         "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-        "    else:\n",
-        "        # We try to install PyTorch3D from source.\n",
-        "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+        "        pip_list = !pip freeze\n",
+        "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+        "    if need_pytorch3d:\n",
+        "        print(f\"failed to find/install wheel for {version_str}\")\n",
+        "if need_pytorch3d:\n",
+        "    print(\"Installing PyTorch3D from source\")\n",
+        "    !pip install ninja\n",
+        "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
       ]
     },
     {

docs/tutorials/render_colored_points.ipynb

@@ -44,25 +44,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/render_densepose.ipynb

@@ -51,25 +51,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/render_textured_meshes.ipynb

@@ -67,25 +67,31 @@
     "import os\n",
     "import sys\n",
     "import torch\n",
+    "import subprocess\n",
     "need_pytorch3d=False\n",
     "try:\n",
     "    import pytorch3d\n",
     "except ModuleNotFoundError:\n",
     "    need_pytorch3d=True\n",
     "if need_pytorch3d:\n",
-    "    if torch.__version__.startswith((\"1.13.\", \"2.0.\")) and sys.platform.startswith(\"linux\"):\n",
-    "        # We try to install PyTorch3D via a released wheel.\n",
-    "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
-    "        version_str=\"\".join([\n",
-    "            f\"py3{sys.version_info.minor}_cu\",\n",
-    "            torch.version.cuda.replace(\".\",\"\"),\n",
-    "            f\"_pyt{pyt_version_str}\"\n",
-    "        ])\n",
-    "        !pip install fvcore iopath\n",
+    "    pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+    "    version_str=\"\".join([\n",
+    "        f\"py3{sys.version_info.minor}_cu\",\n",
+    "        torch.version.cuda.replace(\".\",\"\"),\n",
+    "        f\"_pyt{pyt_version_str}\"\n",
+    "    ])\n",
+    "    !pip install fvcore iopath\n",
+    "    if sys.platform.startswith(\"linux\"):\n",
+    "        print(\"Trying to install wheel for PyTorch3D\")\n",
     "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
-    "    else:\n",
-    "        # We try to install PyTorch3D from source.\n",
-    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "        pip_list = !pip freeze\n",
+    "        need_pytorch3d = not any(i.startswith(\"pytorch3d==\") for  i in pip_list)\n",
+    "    if need_pytorch3d:\n",
+    "        print(f\"failed to find/install wheel for {version_str}\")\n",
+    "if need_pytorch3d:\n",
+    "    print(\"Installing PyTorch3D from source\")\n",
+    "    !pip install ninja\n",
+    "    !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {

docs/tutorials/utils/camera_visualization.py

@@ -33,7 +33,7 @@ def plot_camera_scene(cameras, cameras_gt, status: str):
     a string passed inside the `status` argument.
     """
     fig = plt.figure()
-    ax = fig.gca(projection="3d")
+    ax = fig.add_subplot(projection="3d")
     ax.clear()
     ax.set_title(status)
     handle_cam = plot_cameras(ax, cameras, color="#FF7D1E")

packaging/build_conda.py

@@ -50,7 +50,6 @@ def setup_cuda():
     os.environ["FORCE_CUDA"] = "1"
 
     basic_nvcc_flags = (
-        "-gencode=arch=compute_35,code=sm_35 "
         "-gencode=arch=compute_50,code=sm_50 "
         "-gencode=arch=compute_60,code=sm_60 "
         "-gencode=arch=compute_70,code=sm_70 "
@@ -58,13 +57,19 @@ def setup_cuda():
         "-gencode=arch=compute_50,code=compute_50"
     )
     if CU_VERSION == "cu102":
-        nvcc_flags = basic_nvcc_flags
-    elif CU_VERSION == "cu110":
-        nvcc_flags = "-gencode=arch=compute_80,code=sm_80 " + basic_nvcc_flags
+        nvcc_flags = "-gencode=arch=compute_35,code=sm_35 " + basic_nvcc_flags
+    elif CU_VERSION < ("cu118"):
+        nvcc_flags = (
+            "-gencode=arch=compute_35,code=sm_35 "
+            + "-gencode=arch=compute_80,code=sm_80 "
+            + "-gencode=arch=compute_86,code=sm_86 "
+            + basic_nvcc_flags
+        )
     else:
         nvcc_flags = (
             "-gencode=arch=compute_80,code=sm_80 "
             + "-gencode=arch=compute_86,code=sm_86 "
+            + "-gencode=arch=compute_90,code=sm_90 "
             + basic_nvcc_flags
         )
 
@@ -75,6 +80,12 @@ def setup_cuda():
 def setup_conda_pytorch_constraint() -> List[str]:
     pytorch_constraint = f"- pytorch=={PYTORCH_VERSION}"
     os.environ["CONDA_PYTORCH_CONSTRAINT"] = pytorch_constraint
+    if pytorch_major_minor < (2, 2):
+        os.environ["CONDA_PYTORCH_MKL_CONSTRAINT"] = "- mkl!=2024.1.0"
+        os.environ["SETUPTOOLS_CONSTRAINT"] = "- setuptools<70"
+    else:
+        os.environ["CONDA_PYTORCH_MKL_CONSTRAINT"] = ""
+        os.environ["SETUPTOOLS_CONSTRAINT"] = "- setuptools"
     os.environ["CONDA_PYTORCH_BUILD_CONSTRAINT"] = pytorch_constraint
     os.environ["PYTORCH_VERSION_NODOT"] = PYTORCH_VERSION.replace(".", "")
 

packaging/linux_wheels/go.sh

@@ -5,7 +5,13 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-sudo docker run --rm  -v "$PWD/../../:/inside" pytorch/conda-cuda bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm  -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu113 pytorch/conda-builder:cuda113 bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm  -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu115 pytorch/conda-builder:cuda115 bash inside/packaging/linux_wheels/inside.sh
-sudo docker run --rm  -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu116 pytorch/conda-builder:cuda116 bash inside/packaging/linux_wheels/inside.sh
+# Some directory to persist downloaded conda packages
+conda_cache=/raid/$USER/building_conda_cache
+
+mkdir -p "$conda_cache"
+
+sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu113 pytorch/conda-builder:cuda113 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu115 pytorch/conda-builder:cuda115 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu116 pytorch/conda-builder:cuda116 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu117 pytorch/conda-builder:cuda117 bash inside/packaging/linux_wheels/inside.sh
+sudo docker run --rm -v "$conda_cache:/conda_cache" -v "$PWD/../../:/inside" -e SELECTED_CUDA=cu118 pytorch/conda-builder:cuda118 bash inside/packaging/linux_wheels/inside.sh

packaging/linux_wheels/inside.sh

@@ -16,23 +16,32 @@ VERSION=$(python -c "exec(open('pytorch3d/__init__.py').read()); print(__version
 
 export BUILD_VERSION=$VERSION
 export FORCE_CUDA=1
+export MAX_JOBS=8
+export CONDA_PKGS_DIRS=/conda_cache
 
-wget --no-verbose https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
-tar xzf 1.10.0.tar.gz
-CUB_HOME=$(realpath ./cub-1.10.0)
-export CUB_HOME
-echo "CUB_HOME is now $CUB_HOME"
+if false
+then
+    # We used to have to do this for old versions of CUDA
+    wget --no-verbose https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
+    tar xzf 1.10.0.tar.gz
+    CUB_HOME=$(realpath ./cub-1.10.0)
+    export CUB_HOME
+    echo "CUB_HOME is now $CUB_HOME"
+fi
 
 # As a rule, we want to build for any combination of dependencies which is supported by
 # PyTorch3D and not older than the current Google Colab set up.
 
-PYTHON_VERSIONS="3.7 3.8 3.9 3.10"
+PYTHON_VERSIONS="3.8 3.9 3.10"
 # the keys are pytorch versions
 declare -A CONDA_CUDA_VERSIONS=(
-    ["1.10.1"]="cu111 cu113"
-    ["1.10.2"]="cu111 cu113"
-    ["1.10.0"]="cu111 cu113"
-    ["1.11.0"]="cu111 cu113 cu115"
+#    ["1.11.0"]="cu113"
+#    ["1.12.0"]="cu113"
+#    ["1.12.1"]="cu113"
+#    ["1.13.0"]="cu116"
+#    ["1.13.1"]="cu116 cu117"
+#    ["2.0.0"]="cu117 cu118"
+    ["2.0.1"]="cu117 cu118"
 )
 
 
@@ -41,39 +50,43 @@ for python_version in $PYTHON_VERSIONS
 do
     for pytorch_version in "${!CONDA_CUDA_VERSIONS[@]}"
     do
-        if [[ "3.7 3.8" != *$python_version* ]] && [[ "1.7.0" == *$pytorch_version* ]]
-        then
-            #python 3.9 and later not supported by pytorch 1.7.0 and before
-            continue
-        fi
         if [[ "3.7 3.8 3.9" != *$python_version* ]] && [[ "1.7.0 1.7.1 1.8.0 1.8.1 1.9.0 1.9.1 1.10.0 1.10.1 1.10.2" == *$pytorch_version* ]]
         then
             #python 3.10 and later not supported by pytorch 1.10.2 and before
             continue
         fi
 
-        extra_channel="-c conda-forge"
+        extra_channel="-c nvidia"
+        cudatools="pytorch-cuda"
         if [[ "1.11.0" == "$pytorch_version" ]]
         then
             extra_channel=""
+            cudatools="cudatoolkit"
+        fi
+        if [[ "1.12.0" == "$pytorch_version" ]] || [[ "1.12.1" == "$pytorch_version" ]]
+        then
+            extra_channel="-c conda-forge"
+            cudatools="cudatoolkit"
         fi
 
         for cu_version in ${CONDA_CUDA_VERSIONS[$pytorch_version]}
         do
-            if [[ "cu113 cu115 cu116" == *$cu_version* ]]
-            #       ^^^ CUDA versions listed here have to be built
-            # in their own containers.
-            then
             if [[ $SELECTED_CUDA != "$cu_version" ]]
-                then
-                    continue
-                fi
-            elif [[ $SELECTED_CUDA != "" ]]
             then
                 continue
             fi
 
             case "$cu_version" in
+                cu118)
+                    export CUDA_HOME=/usr/local/cuda-11.8/
+                    export CUDA_TAG=11.8
+                    export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
+                ;;
+                cu117)
+                    export CUDA_HOME=/usr/local/cuda-11.7/
+                    export CUDA_TAG=11.7
+                    export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
+                ;;
                 cu116)
                     export CUDA_HOME=/usr/local/cuda-11.6/
                     export CUDA_TAG=11.6
@@ -130,7 +143,7 @@ do
             conda create -y -n "$tag" "python=$python_version"
             conda activate "$tag"
             # shellcheck disable=SC2086
-            conda install -y -c pytorch $extra_channel "pytorch=$pytorch_version" "cudatoolkit=$CUDA_TAG" torchvision
+            conda install -y -c pytorch $extra_channel "pytorch=$pytorch_version" "$cudatools=$CUDA_TAG"
             pip install fvcore iopath
             echo "python version" "$python_version" "pytorch version" "$pytorch_version" "cuda version" "$cu_version" "tag" "$tag"
 

packaging/pytorch3d/meta.yaml

@@ -12,8 +12,9 @@ requirements:
 
   host:
     - python
-    - setuptools
+    {{ environ.get('SETUPTOOLS_CONSTRAINT') }}
     {{ environ.get('CONDA_PYTORCH_BUILD_CONSTRAINT') }}
+    {{ environ.get('CONDA_PYTORCH_MKL_CONSTRAINT') }}
     {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
     {{ environ.get('CONDA_CPUONLY_FEATURE') }}
 

projects/implicitron_trainer/__init__.py

@@ -3,3 +3,5 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe

projects/implicitron_trainer/experiment.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 """"
 This file is the entry point for launching experiments with Implicitron.
 

projects/implicitron_trainer/impl/__init__.py

@@ -3,3 +3,5 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe

projects/implicitron_trainer/impl/model_factory.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import logging
 import os
 from typing import Optional

projects/implicitron_trainer/impl/optimizer_factory.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import inspect
 import logging
 import os
@@ -121,7 +123,6 @@ class ImplicitronOptimizerFactory(OptimizerFactoryBase):
         """
         # Get the parameters to optimize
         if hasattr(model, "_get_param_groups"):  # use the model function
-            # pyre-ignore[29]
             p_groups = model._get_param_groups(self.lr, wd=self.weight_decay)
         else:
             p_groups = [

projects/implicitron_trainer/impl/training_loop.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import logging
 import os
 import time
@@ -21,7 +23,6 @@ from pytorch3d.implicitron.tools.config import (
     run_auto_creation,
 )
 from pytorch3d.implicitron.tools.stats import Stats
-from pytorch3d.renderer.cameras import CamerasBase
 from torch.utils.data import DataLoader, Dataset
 
 from .utils import seed_all_random_engines
@@ -111,6 +112,8 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
     def __post_init__(self):
         run_auto_creation(self)
 
+    # pyre-fixme[14]: `run` overrides method defined in `TrainingLoopBase`
+    #  inconsistently.
     def run(
         self,
         *,
@@ -256,7 +259,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
             list(log_vars),
             plot_file=os.path.join(exp_dir, "train_stats.pdf"),
             visdom_env=visdom_env_charts,
-            verbose=False,
             visdom_server=self.visdom_server,
             visdom_port=self.visdom_port,
         )
@@ -382,7 +384,8 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
 
             # print textual status update
             if it % self.metric_print_interval == 0 or last_iter:
-                stats.print(stat_set=trainmode, max_it=n_batches)
+                std_out = stats.get_status_string(stat_set=trainmode, max_it=n_batches)
+                logger.info(std_out)
 
             # visualize results
             if (
@@ -392,7 +395,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
             ):
                 prefix = f"e{stats.epoch}_it{stats.it[trainmode]}"
                 if hasattr(model, "visualize"):
-                    # pyre-ignore [29]
                     model.visualize(
                         viz,
                         visdom_env_imgs,

projects/implicitron_trainer/impl/utils.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import random
 

projects/implicitron_trainer/tests/__init__.py

@@ -3,3 +3,5 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe

projects/implicitron_trainer/tests/experiment.yaml

@@ -129,6 +129,19 @@ data_source_ImplicitronDataSource_args:
     dataset_length_train: 0
     dataset_length_val: 0
     dataset_length_test: 0
+  data_loader_map_provider_TrainEvalDataLoaderMapProvider_args:
+    batch_size: 1
+    num_workers: 0
+    dataset_length_train: 0
+    dataset_length_val: 0
+    dataset_length_test: 0
+    train_conditioning_type: SAME
+    val_conditioning_type: SAME
+    test_conditioning_type: KNOWN
+    images_per_seq_options: []
+    sample_consecutive_frames: false
+    consecutive_frames_max_gap: 0
+    consecutive_frames_max_gap_seconds: 0.1
 model_factory_ImplicitronModelFactory_args:
   resume: true
   model_class_type: GenericModel
@@ -203,6 +216,7 @@ model_factory_ImplicitronModelFactory_args:
       n_rays_total_training: null
       stratified_point_sampling_training: true
       stratified_point_sampling_evaluation: false
+      cast_ray_bundle_as_cone: false
       scene_extent: 8.0
       scene_center:
       - 0.0
@@ -215,6 +229,7 @@ model_factory_ImplicitronModelFactory_args:
       n_rays_total_training: null
       stratified_point_sampling_training: true
       stratified_point_sampling_evaluation: false
+      cast_ray_bundle_as_cone: false
       min_depth: 0.1
       max_depth: 8.0
     renderer_LSTMRenderer_args:
@@ -234,6 +249,8 @@ model_factory_ImplicitronModelFactory_args:
       append_coarse_samples_to_fine: true
       density_noise_std_train: 0.0
       return_weights: false
+      blurpool_weights: false
+      sample_pdf_eps: 1.0e-05
       raymarcher_CumsumRaymarcher_args:
         surface_thickness: 1
         bg_color:
@@ -346,6 +363,7 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
+      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 2.0
       n_hidden_neurons_xyz: 80
       n_layers_xyz: 2
@@ -357,6 +375,7 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
+      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 1.0
       n_hidden_neurons_xyz: 256
       n_layers_xyz: 8
@@ -629,6 +648,7 @@ model_factory_ImplicitronModelFactory_args:
       n_rays_total_training: null
       stratified_point_sampling_training: true
       stratified_point_sampling_evaluation: false
+      cast_ray_bundle_as_cone: false
       scene_extent: 8.0
       scene_center:
       - 0.0
@@ -641,6 +661,7 @@ model_factory_ImplicitronModelFactory_args:
       n_rays_total_training: null
       stratified_point_sampling_training: true
       stratified_point_sampling_evaluation: false
+      cast_ray_bundle_as_cone: false
       min_depth: 0.1
       max_depth: 8.0
     renderer_LSTMRenderer_args:
@@ -660,6 +681,8 @@ model_factory_ImplicitronModelFactory_args:
       append_coarse_samples_to_fine: true
       density_noise_std_train: 0.0
       return_weights: false
+      blurpool_weights: false
+      sample_pdf_eps: 1.0e-05
       raymarcher_CumsumRaymarcher_args:
         surface_thickness: 1
         bg_color:
@@ -724,6 +747,7 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
+      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 2.0
       n_hidden_neurons_xyz: 80
       n_layers_xyz: 2
@@ -735,6 +759,7 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
+      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 1.0
       n_hidden_neurons_xyz: 256
       n_layers_xyz: 8
@@ -962,6 +987,7 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
+      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 2.0
       n_hidden_neurons_xyz: 80
       n_layers_xyz: 2
@@ -973,6 +999,7 @@ model_factory_ImplicitronModelFactory_args:
       n_hidden_neurons_dir: 128
       input_xyz: true
       xyz_ray_dir_in_camera_coords: false
+      use_integrated_positional_encoding: false
       transformer_dim_down_factor: 1.0
       n_hidden_neurons_xyz: 256
       n_layers_xyz: 8

projects/implicitron_trainer/tests/test_experiment.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import os
 import tempfile
 import unittest
@@ -132,6 +134,13 @@ class TestExperiment(unittest.TestCase):
         # Check that the default config values, defined by Experiment and its
         # members, is what we expect it to be.
         cfg = OmegaConf.structured(experiment.Experiment)
+        # the following removes the possible effect of env variables
+        ds_arg = cfg.data_source_ImplicitronDataSource_args
+        ds_arg.dataset_map_provider_JsonIndexDatasetMapProvider_args.dataset_root = ""
+        ds_arg.dataset_map_provider_JsonIndexDatasetMapProviderV2_args.dataset_root = ""
+        if "dataset_map_provider_SqlIndexDatasetMapProvider_args" in ds_arg:
+            del ds_arg.dataset_map_provider_SqlIndexDatasetMapProvider_args
+        cfg.training_loop_ImplicitronTrainingLoop_args.visdom_port = 8097
         yaml = OmegaConf.to_yaml(cfg, sort_keys=False)
         if DEBUG:
             (DATA_DIR / "experiment.yaml").write_text(yaml)

projects/implicitron_trainer/tests/test_optimizer_factory.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import logging
 import os
 import unittest

projects/implicitron_trainer/tests/test_visualize.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import os
 import unittest
 

projects/implicitron_trainer/tests/utils.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import contextlib
 import logging
 import os

projects/implicitron_trainer/visualize_reconstruction.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 """
 Script to visualize a previously trained model. Example call:
 

projects/nerf/nerf/nerf_renderer.py

@@ -343,12 +343,14 @@ class RadianceFieldRenderer(torch.nn.Module):
             # For a full render pass concatenate the output chunks,
             # and reshape to image size.
             out = {
-                k: torch.cat(
-                    [ch_o[k] for ch_o in chunk_outputs],
-                    dim=1,
-                ).view(-1, *self._image_size, 3)
-                if chunk_outputs[0][k] is not None
-                else None
+                k: (
+                    torch.cat(
+                        [ch_o[k] for ch_o in chunk_outputs],
+                        dim=1,
+                    ).view(-1, *self._image_size, 3)
+                    if chunk_outputs[0][k] is not None
+                    else None
+                )
                 for k in ("rgb_fine", "rgb_coarse", "rgb_gt")
             }
         else:

projects/nerf/nerf/raysampler.py

@@ -330,9 +330,9 @@ class NeRFRaysampler(torch.nn.Module):
 
             if self.training:
                 # During training we randomly subsample rays.
-                sel_rays = torch.randperm(n_pixels, device=device)[
-                    : self._mc_raysampler._n_rays_per_image
-                ]
+                sel_rays = torch.randperm(
+                    n_pixels, device=full_ray_bundle.lengths.device
+                )[: self._mc_raysampler._n_rays_per_image]
             else:
                 # In case we test, we take only the requested chunk.
                 if chunksize is None:

pytorch3d/__init__.py

@@ -4,4 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-__version__ = "0.7.3"
+# pyre-unsafe
+
+__version__ = "0.7.6"

pytorch3d/common/__init__.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from .datatypes import Device, get_device, make_device
 
 

pytorch3d/common/compat.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from typing import Sequence, Tuple, Union
 
 import torch

pytorch3d/common/datatypes.py

@@ -4,7 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import sys
+# pyre-unsafe
+
 from typing import Optional, Union
 
 import torch

pytorch3d/common/linear_with_repeat.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import math
 from typing import Tuple
 

pytorch3d/common/workaround/__init__.py

@@ -4,5 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from .symeig3x3 import symeig3x3
 from .utils import _safe_det_3x3

pytorch3d/common/workaround/symeig3x3.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import math
 from typing import Optional, Tuple
 

pytorch3d/common/workaround/utils.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import torch
 

pytorch3d/csrc/face_areas_normals/face_areas_normals.cu

@@ -266,6 +266,8 @@ at::Tensor FaceAreasNormalsBackwardCuda(
       grad_normals_t{grad_normals, "grad_normals", 4};
   at::CheckedFrom c = "FaceAreasNormalsBackwardCuda";
   at::checkAllSameGPU(c, {verts_t, faces_t, grad_areas_t, grad_normals_t});
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("FaceAreasNormalsBackwardCuda");
 
   // Set the device for the kernel launch based on the device of verts
   at::cuda::CUDAGuard device_guard(verts.device());

pytorch3d/csrc/interp_face_attrs/interp_face_attrs.cu

@@ -130,6 +130,9 @@ std::tuple<at::Tensor, at::Tensor> InterpFaceAttrsBackwardCuda(
   at::checkAllSameType(
       c, {barycentric_coords_t, face_attrs_t, grad_pix_attrs_t});
 
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("InterpFaceAttrsBackwardCuda");
+
   // Set the device for the kernel launch based on the input
   at::cuda::CUDAGuard device_guard(pix_to_face.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/iou_box3d/iou_box3d.cu

@@ -12,8 +12,6 @@
 #include <math.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <thrust/device_vector.h>
-#include <thrust/tuple.h>
 #include "iou_box3d/iou_utils.cuh"
 
 // Parallelize over N*M computations which can each be done

pytorch3d/csrc/iou_box3d/iou_utils.cuh

@@ -8,7 +8,6 @@
 
 #include <float.h>
 #include <math.h>
-#include <thrust/device_vector.h>
 #include <cstdio>
 #include "utils/float_math.cuh"
 

pytorch3d/csrc/knn/knn.cu

@@ -338,7 +338,7 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborIdxCuda(
 
   TORCH_CHECK((norm == 1) || (norm == 2), "Norm must be 1 or 2.");
 
-  TORCH_CHECK(p2.size(2) == D, "Point sets must have the same last dimension");
+  TORCH_CHECK(p1.size(2) == D, "Point sets must have the same last dimension");
   auto long_dtype = lengths1.options().dtype(at::kLong);
   auto idxs = at::zeros({N, P1, K}, long_dtype);
   auto dists = at::zeros({N, P1, K}, p1.options());
@@ -495,7 +495,7 @@ __global__ void KNearestNeighborBackwardKernel(
     if ((p1_idx < num1) && (k < num2)) {
       const float grad_dist = grad_dists[n * P1 * K + p1_idx * K + k];
       // index of point in p2 corresponding to the k-th nearest neighbor
-      const size_t p2_idx = idxs[n * P1 * K + p1_idx * K + k];
+      const int64_t p2_idx = idxs[n * P1 * K + p1_idx * K + k];
       // If the index is the pad value of -1 then ignore it
       if (p2_idx == -1) {
         continue;
@@ -534,6 +534,9 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborBackwardCuda(
       c, {p1_t, p2_t, lengths1_t, lengths2_t, idxs_t, grad_dists_t});
   at::checkAllSameType(c, {p1_t, p2_t, grad_dists_t});
 
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("KNearestNeighborBackwardCuda");
+
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(p1.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/marching_cubes/marching_cubes.cu

@@ -9,8 +9,6 @@
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
-#include <thrust/device_vector.h>
-#include <thrust/scan.h>
 #include <cstdio>
 #include "marching_cubes/tables.h"
 
@@ -40,20 +38,6 @@ through" each cube in the grid.
 // EPS: Used to indicate if two float values are close
 __constant__ const float EPSILON = 1e-5;
 
-// Thrust wrapper for exclusive scan
-//
-// Args:
-//    output: pointer to on-device output array
-//    input: pointer to on-device input array, where scan is performed
-//    numElements: number of elements for the input array
-//
-void ThrustScanWrapper(int* output, int* input, int numElements) {
-  thrust::exclusive_scan(
-      thrust::device_ptr<int>(input),
-      thrust::device_ptr<int>(input + numElements),
-      thrust::device_ptr<int>(output));
-}
-
 // Linearly interpolate the position where an isosurface cuts an edge
 // between two vertices, based on their scalar values
 //
@@ -239,7 +223,7 @@ __global__ void CompactVoxelsKernel(
         compactedVoxelArray,
     const at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits>
         voxelOccupied,
-    const at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits>
+    const at::PackedTensorAccessor32<int64_t, 1, at::RestrictPtrTraits>
         voxelOccupiedScan,
     uint numVoxels) {
   uint id = blockIdx.x * blockDim.x + threadIdx.x;
@@ -271,7 +255,8 @@ __global__ void GenerateFacesKernel(
     at::PackedTensorAccessor<int64_t, 1, at::RestrictPtrTraits> ids,
     at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits>
         compactedVoxelArray,
-    at::PackedTensorAccessor32<int, 1, at::RestrictPtrTraits> numVertsScanned,
+    at::PackedTensorAccessor32<int64_t, 1, at::RestrictPtrTraits>
+        numVertsScanned,
     const uint activeVoxels,
     const at::PackedTensorAccessor32<float, 3, at::RestrictPtrTraits> vol,
     const at::PackedTensorAccessor32<int, 2, at::RestrictPtrTraits> faceTable,
@@ -397,6 +382,44 @@ __global__ void GenerateFacesKernel(
   } // end for grid-strided kernel
 }
 
+// ATen/Torch does not have an exclusive-scan operator. Additionally, in the
+// code below we need to get the "total number of items to work on" after
+// a scan, which with an inclusive-scan would simply be the value of the last
+// element in the tensor.
+//
+// This utility function hits two birds with one stone, by running
+// an inclusive-scan into a right-shifted view of a tensor that's
+// allocated to be one element bigger than the input tensor.
+//
+// Note; return tensor is `int64_t` per element, even if the input
+// tensor is only 32-bit. Also, the return tensor is one element bigger
+// than the input one.
+//
+// Secondary optional argument is an output argument that gets the
+// value of the last element of the return tensor (because you almost
+// always need this CPU-side right after this function anyway).
+static at::Tensor ExclusiveScanAndTotal(
+    const at::Tensor& inTensor,
+    int64_t* optTotal = nullptr) {
+  const auto inSize = inTensor.sizes()[0];
+  auto retTensor = at::zeros({inSize + 1}, at::kLong).to(inTensor.device());
+
+  using at::indexing::None;
+  using at::indexing::Slice;
+  auto rightShiftedView = retTensor.index({Slice(1, None)});
+
+  // Do an (inclusive-scan) cumulative sum in to the view that's
+  // shifted one element to the right...
+  at::cumsum_out(rightShiftedView, inTensor, 0, at::kLong);
+
+  if (optTotal) {
+    *optTotal = retTensor[inSize].cpu().item<int64_t>();
+  }
+
+  // ...so that the not-shifted tensor holds the exclusive-scan
+  return retTensor;
+}
+
 // Entrance for marching cubes cuda extension. Marching Cubes is an algorithm to
 // create triangle meshes from an implicit function (one of the form f(x, y, z)
 // = 0). It works by iteratively checking a grid of cubes superimposed over a
@@ -455,6 +478,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
     grid.x = 65535;
   }
 
+  using at::indexing::None;
+  using at::indexing::Slice;
+
   auto d_voxelVerts =
       at::zeros({numVoxels}, at::TensorOptions().dtype(at::kInt))
           .to(vol.device());
@@ -477,18 +503,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
   // count for voxels in the grid and compute the number of active voxels.
   // If the number of active voxels is 0, return zero tensor for verts and
   // faces.
+  int64_t activeVoxels = 0;
   auto d_voxelOccupiedScan =
-      at::zeros({numVoxels}, at::TensorOptions().dtype(at::kInt))
-          .to(vol.device());
-  ThrustScanWrapper(
-      d_voxelOccupiedScan.data_ptr<int>(),
-      d_voxelOccupied.data_ptr<int>(),
-      numVoxels);
-
-  // number of active voxels
-  int lastElement = d_voxelVerts[numVoxels - 1].cpu().item<int>();
-  int lastScan = d_voxelOccupiedScan[numVoxels - 1].cpu().item<int>();
-  int activeVoxels = lastElement + lastScan;
+      ExclusiveScanAndTotal(d_voxelOccupied, &activeVoxels);
 
   const int device_id = vol.device().index();
   auto opt = at::TensorOptions().dtype(at::kInt).device(at::kCUDA, device_id);
@@ -503,28 +520,21 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
     return std::make_tuple(verts, faces, ids);
   }
 
-  // Execute "CompactVoxelsKernel" kernel to compress voxels for accleration.
+  // Execute "CompactVoxelsKernel" kernel to compress voxels for acceleration.
   // This allows us to run triangle generation on only the occupied voxels.
   auto d_compVoxelArray = at::zeros({activeVoxels}, opt);
   CompactVoxelsKernel<<<grid, threads, 0, stream>>>(
       d_compVoxelArray.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
       d_voxelOccupied.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
-      d_voxelOccupiedScan.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
+      d_voxelOccupiedScan
+          .packed_accessor32<int64_t, 1, at::RestrictPtrTraits>(),
       numVoxels);
   AT_CUDA_CHECK(cudaGetLastError());
   cudaDeviceSynchronize();
 
   // Scan d_voxelVerts array to generate offsets of vertices for each voxel
-  auto d_voxelVertsScan = at::zeros({numVoxels}, opt);
-  ThrustScanWrapper(
-      d_voxelVertsScan.data_ptr<int>(),
-      d_voxelVerts.data_ptr<int>(),
-      numVoxels);
-
-  // total number of vertices
-  lastElement = d_voxelVerts[numVoxels - 1].cpu().item<int>();
-  lastScan = d_voxelVertsScan[numVoxels - 1].cpu().item<int>();
-  int totalVerts = lastElement + lastScan;
+  int64_t totalVerts = 0;
+  auto d_voxelVertsScan = ExclusiveScanAndTotal(d_voxelVerts, &totalVerts);
 
   // Execute "GenerateFacesKernel" kernel
   // This runs only on the occupied voxels.
@@ -544,7 +554,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
       faces.packed_accessor<int64_t, 2, at::RestrictPtrTraits>(),
       ids.packed_accessor<int64_t, 1, at::RestrictPtrTraits>(),
       d_compVoxelArray.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
-      d_voxelVertsScan.packed_accessor32<int, 1, at::RestrictPtrTraits>(),
+      d_voxelVertsScan.packed_accessor32<int64_t, 1, at::RestrictPtrTraits>(),
       activeVoxels,
       vol.packed_accessor32<float, 3, at::RestrictPtrTraits>(),
       faceTable.packed_accessor32<int, 2, at::RestrictPtrTraits>(),

pytorch3d/csrc/marching_cubes/marching_cubes_cpu.cpp

@@ -71,8 +71,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCpu(
           if ((j + 1) % 3 == 0 && ps[0] != ps[1] && ps[1] != ps[2] &&
               ps[2] != ps[0]) {
             for (int k = 0; k < 3; k++) {
-              int v = tri[k];
-              edge_id_to_v[tri.at(k)] = ps.at(k);
+              int64_t v = tri.at(k);
+              edge_id_to_v[v] = ps.at(k);
               if (!uniq_edge_id.count(v)) {
                 uniq_edge_id[v] = verts.size();
                 verts.push_back(edge_id_to_v[v]);

pytorch3d/csrc/point_mesh/point_mesh_cuda.cu

@@ -305,6 +305,8 @@ std::tuple<at::Tensor, at::Tensor> DistanceBackwardCuda(
   at::CheckedFrom c = "DistanceBackwardCuda";
   at::checkAllSameGPU(c, {objects_t, targets_t, idx_objects_t, grad_dists_t});
   at::checkAllSameType(c, {objects_t, targets_t, grad_dists_t});
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("DistanceBackwardCuda");
 
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(objects.device());
@@ -624,6 +626,9 @@ std::tuple<at::Tensor, at::Tensor> PointFaceArrayDistanceBackwardCuda(
   at::CheckedFrom c = "PointFaceArrayDistanceBackwardCuda";
   at::checkAllSameGPU(c, {points_t, tris_t, grad_dists_t});
   at::checkAllSameType(c, {points_t, tris_t, grad_dists_t});
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic(
+      "PointFaceArrayDistanceBackwardCuda");
 
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points.device());
@@ -787,6 +792,9 @@ std::tuple<at::Tensor, at::Tensor> PointEdgeArrayDistanceBackwardCuda(
   at::CheckedFrom c = "PointEdgeArrayDistanceBackwardCuda";
   at::checkAllSameGPU(c, {points_t, segms_t, grad_dists_t});
   at::checkAllSameType(c, {points_t, segms_t, grad_dists_t});
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic(
+      "PointEdgeArrayDistanceBackwardCuda");
 
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points.device());

pytorch3d/csrc/points_to_volumes/points_to_volumes.cu

@@ -141,6 +141,9 @@ void PointsToVolumesForwardCuda(
        grid_sizes_t,
        mask_t});
 
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("PointsToVolumesForwardCuda");
+
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points_3d.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/pulsar/global.h

@@ -30,11 +30,18 @@
 #define GLOBAL __global__
 #define RESTRICT __restrict__
 #define DEBUGBREAK()
+#ifdef __NVCC_DIAG_PRAGMA_SUPPORT__
+#pragma nv_diag_suppress 1866
+#pragma nv_diag_suppress 2941
+#pragma nv_diag_suppress 2951
+#pragma nv_diag_suppress 2967
+#else
 #pragma diag_suppress = attribute_not_allowed
 #pragma diag_suppress = 1866
 #pragma diag_suppress = 2941
 #pragma diag_suppress = 2951
 #pragma diag_suppress = 2967
+#endif
 #else // __CUDACC__
 #define INLINE inline
 #define HOST
@@ -49,6 +56,7 @@
 #pragma clang diagnostic pop
 #ifdef WITH_CUDA
 #include <ATen/cuda/CUDAContext.h>
+#include <vector_functions.h>
 #else
 #ifndef cudaStream_t
 typedef void* cudaStream_t;
@@ -65,8 +73,6 @@ struct float2 {
 struct float3 {
   float x, y, z;
 };
-#endif
-namespace py = pybind11;
 inline float3 make_float3(const float& x, const float& y, const float& z) {
   float3 res;
   res.x = x;
@@ -74,6 +80,8 @@ inline float3 make_float3(const float& x, const float& y, const float& z) {
   res.z = z;
   return res;
 }
+#endif
+namespace py = pybind11;
 
 inline bool operator==(const float3& a, const float3& b) {
   return a.x == b.x && a.y == b.y && a.z == b.z;

pytorch3d/csrc/pulsar/host/commands.h

@@ -357,11 +357,11 @@ void MAX_WS(
 //
 //
 #define END_PARALLEL() \
-  end_parallel:;       \
+  end_parallel :;      \
   }
 #define END_PARALLEL_NORET() }
 #define END_PARALLEL_2D() \
-  end_parallel:;          \
+  end_parallel :;         \
   }                       \
   }
 #define END_PARALLEL_2D_NORET() \

pytorch3d/csrc/pulsar/include/renderer.construct.device.h

@@ -93,7 +93,7 @@ HOST void construct(
   MALLOC(self->di_sorted_d, DrawInfo, max_num_balls);
   MALLOC(self->region_flags_d, char, max_num_balls);
   MALLOC(self->num_selected_d, size_t, 1);
-  MALLOC(self->forw_info_d, float, width* height*(3 + 2 * n_track));
+  MALLOC(self->forw_info_d, float, width* height * (3 + 2 * n_track));
   MALLOC(self->min_max_pixels_d, IntersectInfo, 1);
   MALLOC(self->grad_pos_d, float3, max_num_balls);
   MALLOC(self->grad_col_d, float, max_num_balls* n_channels);

pytorch3d/csrc/pulsar/include/renderer.forward.device.h

@@ -102,6 +102,7 @@ void forward(
       self->workspace_d,
       self->workspace_size,
       stream);
+  CHECKLAUNCH();
   SORT_ASCENDING_WS(
       self->min_depth_d,
       self->min_depth_sorted_d,
@@ -111,6 +112,7 @@ void forward(
       self->workspace_d,
       self->workspace_size,
       stream);
+  CHECKLAUNCH();
   SORT_ASCENDING_WS(
       self->min_depth_d,
       self->min_depth_sorted_d,

pytorch3d/csrc/pulsar/include/renderer.render.device.h

@@ -99,7 +99,7 @@ GLOBAL void render(
   /** Whether loading of balls is completed. */
   SHARED bool loading_done;
   /** The number of balls loaded overall (just for statistics). */
-  SHARED int n_balls_loaded;
+  [[maybe_unused]] SHARED int n_balls_loaded;
   /** The area this thread block covers. */
   SHARED IntersectInfo block_area;
   if (thread_block.thread_rank() == 0) {

pytorch3d/csrc/pulsar/pytorch/camera.h

@@ -37,7 +37,7 @@ inline void fill_cam_vecs(
   res->pixel_dir_y.x = pixel_dir_y.data_ptr<float>()[0];
   res->pixel_dir_y.y = pixel_dir_y.data_ptr<float>()[1];
   res->pixel_dir_y.z = pixel_dir_y.data_ptr<float>()[2];
-  auto sensor_dir_z = pixel_dir_y.cross(pixel_dir_x);
+  auto sensor_dir_z = pixel_dir_y.cross(pixel_dir_x, -1);
   sensor_dir_z /= sensor_dir_z.norm();
   if (right_handed) {
     sensor_dir_z *= -1.f;

pytorch3d/csrc/rasterize_coarse/rasterize_coarse.cu

@@ -244,8 +244,7 @@ at::Tensor RasterizeCoarseCuda(
   if (num_bins_y >= kMaxItemsPerBin || num_bins_x >= kMaxItemsPerBin) {
     std::stringstream ss;
     ss << "In RasterizeCoarseCuda got num_bins_y: " << num_bins_y
-       << ", num_bins_x: " << num_bins_x << ", "
-       << "; that's too many!";
+       << ", num_bins_x: " << num_bins_x << ", " << "; that's too many!";
     AT_ERROR(ss.str());
   }
   auto opts = elems_per_batch.options().dtype(at::kInt);

pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu

@@ -583,6 +583,9 @@ at::Tensor RasterizeMeshesBackwardCuda(
   at::checkAllSameType(
       c, {face_verts_t, grad_zbuf_t, grad_bary_t, grad_dists_t});
 
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("RasterizeMeshesBackwardCuda");
+
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(face_verts.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/rasterize_points/rasterize_points.cu

@@ -423,7 +423,8 @@ at::Tensor RasterizePointsBackwardCuda(
   at::CheckedFrom c = "RasterizePointsBackwardCuda";
   at::checkAllSameGPU(c, {points_t, idxs_t, grad_zbuf_t, grad_dists_t});
   at::checkAllSameType(c, {points_t, grad_zbuf_t, grad_dists_t});
-
+  // This is nondeterministic because atomicAdd
+  at::globalContext().alertNotDeterministic("RasterizePointsBackwardCuda");
   // Set the device for the kernel launch based on the device of the input
   at::cuda::CUDAGuard device_guard(points.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();

pytorch3d/csrc/sample_farthest_points/sample_farthest_points.cu

@@ -155,7 +155,7 @@ at::Tensor FarthestPointSamplingCuda(
 
   // Max possible threads per block
   const int MAX_THREADS_PER_BLOCK = 1024;
-  const size_t threads = max(min(1 << points_pow_2, MAX_THREADS_PER_BLOCK), 1);
+  const size_t threads = max(min(1 << points_pow_2, MAX_THREADS_PER_BLOCK), 2);
 
   // Create the accessors
   auto points_a = points.packed_accessor64<float, 3, at::RestrictPtrTraits>();
@@ -215,10 +215,6 @@ at::Tensor FarthestPointSamplingCuda(
       FarthestPointSamplingKernel<2><<<threads, threads, shared_mem, stream>>>(
           points_a, lengths_a, K_a, idxs_a, min_point_dist_a, start_idxs_a);
       break;
-    case 1:
-      FarthestPointSamplingKernel<1><<<threads, threads, shared_mem, stream>>>(
-          points_a, lengths_a, K_a, idxs_a, min_point_dist_a, start_idxs_a);
-      break;
     default:
       FarthestPointSamplingKernel<1024>
           <<<blocks, threads, shared_mem, stream>>>(

pytorch3d/datasets/__init__.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from .r2n2 import BlenderCamera, collate_batched_R2N2, R2N2, render_cubified_voxels
 from .shapenet import ShapeNetCore
 from .utils import collate_batched_meshes

pytorch3d/datasets/r2n2/__init__.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from .r2n2 import R2N2
 from .utils import BlenderCamera, collate_batched_R2N2, render_cubified_voxels
 

pytorch3d/datasets/r2n2/r2n2.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import json
 import warnings
 from os import path

pytorch3d/datasets/r2n2/utils.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import math
 from typing import Dict, List
 

pytorch3d/datasets/shapenet/__init__.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from .shapenet_core import ShapeNetCore
 
 

pytorch3d/datasets/shapenet/shapenet_core.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import json
 import os
 import warnings

pytorch3d/datasets/shapenet_base.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import warnings
 from typing import Dict, List, Optional, Tuple
 

pytorch3d/datasets/utils.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from typing import Dict, List
 
 from pytorch3d.renderer.mesh import TexturesAtlas

pytorch3d/implicitron/__init__.py

@@ -3,3 +3,5 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe

pytorch3d/implicitron/dataset/__init__.py

@@ -3,3 +3,5 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe

pytorch3d/implicitron/dataset/blender_dataset_map_provider.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import torch
 from pytorch3d.implicitron.tools.config import registry

pytorch3d/implicitron/dataset/data_loader_map_provider.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from dataclasses import dataclass
 from enum import Enum
 from typing import Iterator, List, Optional, Tuple

pytorch3d/implicitron/dataset/data_source.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from typing import Optional, Tuple
 
 from pytorch3d.implicitron.tools.config import (
@@ -72,6 +74,16 @@ class ImplicitronDataSource(DataSourceBase):  # pyre-ignore[13]
             from .rendered_mesh_dataset_map_provider import (  # noqa: F401
                 RenderedMeshDatasetMapProvider,
             )
+            from .train_eval_data_loader_provider import (  # noqa: F401
+                TrainEvalDataLoaderMapProvider,
+            )
+
+            try:
+                from .sql_dataset_provider import (  # noqa: F401  # pyre-ignore
+                    SqlIndexDatasetMapProvider,
+                )
+            except ModuleNotFoundError:
+                pass  # environment without SQL dataset
         finally:
             pass
 

pytorch3d/implicitron/dataset/dataset_base.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import (

pytorch3d/implicitron/dataset/dataset_map_provider.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import logging
 import os
 from dataclasses import dataclass

pytorch3d/implicitron/dataset/frame_data.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import os
 from abc import ABC, abstractmethod
 from collections import defaultdict
@@ -203,7 +205,10 @@ class FrameData(Mapping[str, Any]):
                 when no image has been loaded)
         """
         if self.bbox_xywh is None:
-            raise ValueError("Attempted cropping by metadata with empty bounding box")
+            raise ValueError(
+                "Attempted cropping by metadata with empty bounding box. Consider either"
+                " to remove_empty_masks or turn off box_crop in the dataset config."
+            )
 
         if not self._uncropped:
             raise ValueError(
@@ -450,6 +455,9 @@ class FrameDataBuilderBase(ReplaceableBase, Generic[FrameDataSubtype], ABC):
         self,
         frame_annotation: types.FrameAnnotation,
         sequence_annotation: types.SequenceAnnotation,
+        *,
+        load_blobs: bool = True,
+        **kwargs,
     ) -> FrameDataSubtype:
         """An abstract method to build the frame data based on raw frame/sequence
         annotations, load the binary data and adjust them according to the metadata.
@@ -465,8 +473,9 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
     Beware that modifications of frame data are done in-place.
 
     Args:
-        dataset_root: The root folder of the dataset; all the paths in jsons are
-                specified relative to this root (but not json paths themselves).
+        dataset_root: The root folder of the dataset; all paths in frame / sequence
+            annotations are defined w.r.t. this root. Has to be set if any of the
+            load_* flabs below is true.
         load_images: Enable loading the frame RGB data.
         load_depths: Enable loading the frame depth maps.
         load_depth_masks: Enable loading the frame depth map masks denoting the
@@ -494,7 +503,7 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
         path_manager: Optionally a PathManager for interpreting paths in a special way.
     """
 
-    dataset_root: str = ""
+    dataset_root: Optional[str] = None
     load_images: bool = True
     load_depths: bool = True
     load_depth_masks: bool = True
@@ -510,11 +519,32 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
     box_crop_context: float = 0.3
     path_manager: Any = None
 
+    def __post_init__(self) -> None:
+        load_any_blob = (
+            self.load_images
+            or self.load_depths
+            or self.load_depth_masks
+            or self.load_masks
+            or self.load_point_clouds
+        )
+        if load_any_blob and self.dataset_root is None:
+            raise ValueError(
+                "dataset_root must be set to load any blob data. "
+                "Make sure it is set in either FrameDataBuilder or Dataset params."
+            )
+
+        if load_any_blob and not self._exists_in_dataset_root(""):
+            raise ValueError(
+                f"dataset_root is passed but {self.dataset_root} does not exist."
+            )
+
     def build(
         self,
         frame_annotation: types.FrameAnnotation,
         sequence_annotation: types.SequenceAnnotation,
+        *,
         load_blobs: bool = True,
+        **kwargs,
     ) -> FrameDataSubtype:
         """Builds the frame data based on raw frame/sequence annotations, loads the
         binary data and adjust them according to the metadata. The processing includes:
@@ -548,38 +578,58 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
             camera_quality_score=safe_as_tensor(
                 sequence_annotation.viewpoint_quality_score, torch.float
             ),
-            point_cloud_quality_score=safe_as_tensor(
-                point_cloud.quality_score, torch.float
-            )
-            if point_cloud is not None
-            else None,
+            point_cloud_quality_score=(
+                safe_as_tensor(point_cloud.quality_score, torch.float)
+                if point_cloud is not None
+                else None
+            ),
         )
 
-        if load_blobs and self.load_masks and frame_annotation.mask is not None:
-            (
-                frame_data.fg_probability,
-                frame_data.mask_path,
-                frame_data.bbox_xywh,
-            ) = self._load_fg_probability(frame_annotation)
+        fg_mask_np: Optional[np.ndarray] = None
+        mask_annotation = frame_annotation.mask
+        if mask_annotation is not None:
+            if load_blobs and self.load_masks:
+                fg_mask_np, mask_path = self._load_fg_probability(frame_annotation)
+                frame_data.mask_path = mask_path
+                frame_data.fg_probability = safe_as_tensor(fg_mask_np, torch.float)
+
+            bbox_xywh = mask_annotation.bounding_box_xywh
+            if bbox_xywh is None and fg_mask_np is not None:
+                bbox_xywh = get_bbox_from_mask(fg_mask_np, self.box_crop_mask_thr)
+
+            frame_data.bbox_xywh = safe_as_tensor(bbox_xywh, torch.float)
 
         if frame_annotation.image is not None:
             image_size_hw = safe_as_tensor(frame_annotation.image.size, torch.long)
             frame_data.image_size_hw = image_size_hw  # original image size
             # image size after crop/resize
             frame_data.effective_image_size_hw = image_size_hw
+            image_path = None
+            dataset_root = self.dataset_root
+            if frame_annotation.image.path is not None and dataset_root is not None:
+                image_path = os.path.join(dataset_root, frame_annotation.image.path)
+                frame_data.image_path = image_path
 
             if load_blobs and self.load_images:
-                (
-                    frame_data.image_rgb,
-                    frame_data.image_path,
-                ) = self._load_images(frame_annotation, frame_data.fg_probability)
+                if image_path is None:
+                    raise ValueError("Image path is required to load images.")
 
-        if load_blobs and self.load_depths and frame_annotation.depth is not None:
+                image_np = load_image(self._local_path(image_path))
+                frame_data.image_rgb = self._postprocess_image(
+                    image_np, frame_annotation.image.size, frame_data.fg_probability
+                )
+
+        if (
+            load_blobs
+            and self.load_depths
+            and frame_annotation.depth is not None
+            and frame_annotation.depth.path is not None
+        ):
             (
                 frame_data.depth_map,
                 frame_data.depth_path,
                 frame_data.depth_mask,
-            ) = self._load_mask_depth(frame_annotation, frame_data.fg_probability)
+            ) = self._load_mask_depth(frame_annotation, fg_mask_np)
 
         if load_blobs and self.load_point_clouds and point_cloud is not None:
             pcl_path = self._fix_point_cloud_path(point_cloud.path)
@@ -604,66 +654,56 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
 
     def _load_fg_probability(
         self, entry: types.FrameAnnotation
-    ) -> Tuple[Optional[torch.Tensor], Optional[str], Optional[torch.Tensor]]:
-
-        full_path = os.path.join(self.dataset_root, entry.mask.path)  # pyre-ignore
+    ) -> Tuple[np.ndarray, str]:
+        assert self.dataset_root is not None and entry.mask is not None
+        full_path = os.path.join(self.dataset_root, entry.mask.path)
         fg_probability = load_mask(self._local_path(full_path))
-        # we can use provided bbox_xywh or calculate it based on mask
-        # saves time to skip bbox calculation
-        # pyre-ignore
-        bbox_xywh = entry.mask.bounding_box_xywh or get_bbox_from_mask(
-            fg_probability, self.box_crop_mask_thr
-        )
         if fg_probability.shape[-2:] != entry.image.size:
             raise ValueError(
                 f"bad mask size: {fg_probability.shape[-2:]} vs {entry.image.size}!"
             )
-        return (
-            safe_as_tensor(fg_probability, torch.float),
-            full_path,
-            safe_as_tensor(bbox_xywh, torch.long),
-        )
 
-    def _load_images(
+        return fg_probability, full_path
+
+    def _postprocess_image(
         self,
-        entry: types.FrameAnnotation,
+        image_np: np.ndarray,
+        image_size: Tuple[int, int],
         fg_probability: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, str]:
-        assert self.dataset_root is not None and entry.image is not None
-        path = os.path.join(self.dataset_root, entry.image.path)
-        image_rgb = load_image(self._local_path(path))
+    ) -> torch.Tensor:
+        image_rgb = safe_as_tensor(image_np, torch.float)
 
-        if image_rgb.shape[-2:] != entry.image.size:
-            raise ValueError(
-                f"bad image size: {image_rgb.shape[-2:]} vs {entry.image.size}!"
-            )
+        if image_rgb.shape[-2:] != image_size:
+            raise ValueError(f"bad image size: {image_rgb.shape[-2:]} vs {image_size}!")
 
         if self.mask_images:
             assert fg_probability is not None
             image_rgb *= fg_probability
 
-        return image_rgb, path
+        return image_rgb
 
     def _load_mask_depth(
         self,
         entry: types.FrameAnnotation,
-        fg_probability: Optional[torch.Tensor],
+        fg_mask: Optional[np.ndarray],
     ) -> Tuple[torch.Tensor, str, torch.Tensor]:
         entry_depth = entry.depth
-        assert entry_depth is not None
-        path = os.path.join(self.dataset_root, entry_depth.path)
+        dataset_root = self.dataset_root
+        assert dataset_root is not None
+        assert entry_depth is not None and entry_depth.path is not None
+        path = os.path.join(dataset_root, entry_depth.path)
         depth_map = load_depth(self._local_path(path), entry_depth.scale_adjustment)
 
         if self.mask_depths:
-            assert fg_probability is not None
-            depth_map *= fg_probability
+            assert fg_mask is not None
+            depth_map *= fg_mask
 
-        if self.load_depth_masks:
-            assert entry_depth.mask_path is not None
-            mask_path = os.path.join(self.dataset_root, entry_depth.mask_path)
+        mask_path = entry_depth.mask_path
+        if self.load_depth_masks and mask_path is not None:
+            mask_path = os.path.join(dataset_root, mask_path)
             depth_mask = load_depth_mask(self._local_path(mask_path))
         else:
-            depth_mask = torch.ones_like(depth_map)
+            depth_mask = (depth_map > 0.0).astype(np.float32)
 
         return torch.tensor(depth_map), path, torch.tensor(depth_mask)
 
@@ -708,6 +748,7 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
         )
         if path.startswith(unwanted_prefix):
             path = path[len(unwanted_prefix) :]
+        assert self.dataset_root is not None
         return os.path.join(self.dataset_root, path)
 
     def _local_path(self, path: str) -> str:
@@ -715,6 +756,16 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
             return path
         return self.path_manager.get_local_path(path)
 
+    def _exists_in_dataset_root(self, relpath) -> bool:
+        if not self.dataset_root:
+            return False
+
+        full_path = os.path.join(self.dataset_root, relpath)
+        if self.path_manager is None:
+            return os.path.exists(full_path)
+        else:
+            return self.path_manager.exists(full_path)
+
 
 @registry.register
 class FrameDataBuilder(GenericWorkaround, GenericFrameDataBuilder[FrameData]):

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import copy
 import functools
 import gzip
@@ -124,9 +126,9 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             dimension of the cropping bounding box, relative to box size.
     """
 
-    frame_annotations_type: ClassVar[
-        Type[types.FrameAnnotation]
-    ] = types.FrameAnnotation
+    frame_annotations_type: ClassVar[Type[types.FrameAnnotation]] = (
+        types.FrameAnnotation
+    )
 
     path_manager: Any = None
     frame_annotations_file: str = ""
@@ -190,6 +192,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             box_crop=self.box_crop,
             box_crop_mask_thr=self.box_crop_mask_thr,
             box_crop_context=self.box_crop_context,
+            path_manager=self.path_manager,
         )
         logger.info(str(self))
 

pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import json
 import os

pytorch3d/implicitron/dataset/json_index_dataset_map_provider_v2.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import copy
 import json

pytorch3d/implicitron/dataset/llff_dataset_map_provider.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import numpy as np
 import torch

pytorch3d/implicitron/dataset/load_blender.py

@@ -1,6 +1,8 @@
 # @lint-ignore-every LICENSELINT
 # Adapted from https://github.com/bmild/nerf/blob/master/load_blender.py
 # Copyright (c) 2020 bmild
+
+# pyre-unsafe
 import json
 import os
 

pytorch3d/implicitron/dataset/load_llff.py

@@ -1,6 +1,8 @@
 # @lint-ignore-every LICENSELINT
 # Adapted from https://github.com/bmild/nerf/blob/master/load_llff.py
 # Copyright (c) 2020 bmild
+
+# pyre-unsafe
 import logging
 import os
 import warnings
@@ -34,11 +36,7 @@ def _minify(basedir, path_manager, factors=(), resolutions=()):
 
     imgdir = os.path.join(basedir, "images")
     imgs = [os.path.join(imgdir, f) for f in sorted(_ls(path_manager, imgdir))]
-    imgs = [
-        f
-        for f in imgs
-        if any([f.endswith(ex) for ex in ["JPG", "jpg", "png", "jpeg", "PNG"]])
-    ]
+    imgs = [f for f in imgs if f.endswith("JPG", "jpg", "png", "jpeg", "PNG")]
     imgdir_orig = imgdir
 
     wd = os.getcwd()

pytorch3d/implicitron/dataset/orm_types.py

@@ -0,0 +1,189 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This functionality requires SQLAlchemy 2.0 or later.
+
+import math
+import struct
+from typing import Optional, Tuple
+
+import numpy as np
+
+from pytorch3d.implicitron.dataset.types import (
+    DepthAnnotation,
+    ImageAnnotation,
+    MaskAnnotation,
+    PointCloudAnnotation,
+    VideoAnnotation,
+    ViewpointAnnotation,
+)
+
+from sqlalchemy import LargeBinary
+from sqlalchemy.orm import (
+    composite,
+    DeclarativeBase,
+    Mapped,
+    mapped_column,
+    MappedAsDataclass,
+)
+from sqlalchemy.types import TypeDecorator
+
+
+# these produce policies to serialize structured types to blobs
+def ArrayTypeFactory(shape=None):
+    if shape is None:
+
+        class VariableShapeNumpyArrayType(TypeDecorator):
+            impl = LargeBinary
+
+            def process_bind_param(self, value, dialect):
+                if value is None:
+                    return None
+
+                ndim_bytes = np.int32(value.ndim).tobytes()
+                shape_bytes = np.array(value.shape, dtype=np.int64).tobytes()
+                value_bytes = value.astype(np.float32).tobytes()
+                return ndim_bytes + shape_bytes + value_bytes
+
+            def process_result_value(self, value, dialect):
+                if value is None:
+                    return None
+
+                ndim = np.frombuffer(value[:4], dtype=np.int32)[0]
+                value_start = 4 + 8 * ndim
+                shape = np.frombuffer(value[4:value_start], dtype=np.int64)
+                assert shape.shape == (ndim,)
+                return np.frombuffer(value[value_start:], dtype=np.float32).reshape(
+                    shape
+                )
+
+        return VariableShapeNumpyArrayType
+
+    class NumpyArrayType(TypeDecorator):
+        impl = LargeBinary
+
+        def process_bind_param(self, value, dialect):
+            if value is not None:
+                if value.shape != shape:
+                    raise ValueError(f"Passed an array of wrong shape: {value.shape}")
+                return value.astype(np.float32).tobytes()
+            return None
+
+        def process_result_value(self, value, dialect):
+            if value is not None:
+                return np.frombuffer(value, dtype=np.float32).reshape(shape)
+            return None
+
+    return NumpyArrayType
+
+
+def TupleTypeFactory(dtype=float, shape: Tuple[int, ...] = (2,)):
+    format_symbol = {
+        float: "f",  # float32
+        int: "i",  # int32
+    }[dtype]
+
+    class TupleType(TypeDecorator):
+        impl = LargeBinary
+        _format = format_symbol * math.prod(shape)
+
+        def process_bind_param(self, value, _):
+            if value is None:
+                return None
+
+            if len(shape) > 1:
+                value = np.array(value, dtype=dtype).reshape(-1)
+
+            return struct.pack(TupleType._format, *value)
+
+        def process_result_value(self, value, _):
+            if value is None:
+                return None
+
+            loaded = struct.unpack(TupleType._format, value)
+            if len(shape) > 1:
+                loaded = _rec_totuple(
+                    np.array(loaded, dtype=dtype).reshape(shape).tolist()
+                )
+
+            return loaded
+
+    return TupleType
+
+
+def _rec_totuple(t):
+    if isinstance(t, list):
+        return tuple(_rec_totuple(x) for x in t)
+
+    return t
+
+
+class Base(MappedAsDataclass, DeclarativeBase):
+    """subclasses will be converted to dataclasses"""
+
+
+class SqlFrameAnnotation(Base):
+    __tablename__ = "frame_annots"
+
+    sequence_name: Mapped[str] = mapped_column(primary_key=True)
+    frame_number: Mapped[int] = mapped_column(primary_key=True)
+    frame_timestamp: Mapped[float] = mapped_column(index=True)
+
+    image: Mapped[ImageAnnotation] = composite(
+        mapped_column("_image_path"),
+        mapped_column("_image_size", TupleTypeFactory(int)),
+    )
+
+    depth: Mapped[DepthAnnotation] = composite(
+        mapped_column("_depth_path", nullable=True),
+        mapped_column("_depth_scale_adjustment", nullable=True),
+        mapped_column("_depth_mask_path", nullable=True),
+    )
+
+    mask: Mapped[MaskAnnotation] = composite(
+        mapped_column("_mask_path", nullable=True),
+        mapped_column("_mask_mass", index=True, nullable=True),
+        mapped_column(
+            "_mask_bounding_box_xywh",
+            TupleTypeFactory(float, shape=(4,)),
+            nullable=True,
+        ),
+    )
+
+    viewpoint: Mapped[ViewpointAnnotation] = composite(
+        mapped_column(
+            "_viewpoint_R", TupleTypeFactory(float, shape=(3, 3)), nullable=True
+        ),
+        mapped_column(
+            "_viewpoint_T", TupleTypeFactory(float, shape=(3,)), nullable=True
+        ),
+        mapped_column(
+            "_viewpoint_focal_length", TupleTypeFactory(float), nullable=True
+        ),
+        mapped_column(
+            "_viewpoint_principal_point", TupleTypeFactory(float), nullable=True
+        ),
+        mapped_column("_viewpoint_intrinsics_format", nullable=True),
+    )
+
+
+class SqlSequenceAnnotation(Base):
+    __tablename__ = "sequence_annots"
+
+    sequence_name: Mapped[str] = mapped_column(primary_key=True)
+    category: Mapped[str] = mapped_column(index=True)
+
+    video: Mapped[VideoAnnotation] = composite(
+        mapped_column("_video_path", nullable=True),
+        mapped_column("_video_length", nullable=True),
+    )
+    point_cloud: Mapped[PointCloudAnnotation] = composite(
+        mapped_column("_point_cloud_path", nullable=True),
+        mapped_column("_point_cloud_quality_score", nullable=True),
+        mapped_column("_point_cloud_n_points", nullable=True),
+    )
+    # the bigger the better
+    viewpoint_quality_score: Mapped[Optional[float]] = mapped_column()

pytorch3d/implicitron/dataset/rendered_mesh_dataset_map_provider.py

@@ -4,15 +4,13 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from os.path import dirname, join, realpath
 from typing import Optional, Tuple
 
 import torch
-from pytorch3d.implicitron.tools.config import (
-    expand_args_fields,
-    registry,
-    run_auto_creation,
-)
+from pytorch3d.implicitron.tools.config import registry, run_auto_creation
 from pytorch3d.io import IO
 from pytorch3d.renderer import (
     AmbientLights,

pytorch3d/implicitron/dataset/scene_batch_sampler.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import warnings
 from collections import Counter

pytorch3d/implicitron/dataset/single_sequence_dataset.py

@@ -4,12 +4,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 # This file defines a base class for dataset map providers which
 # provide data for a single scene.
 
 from dataclasses import field
-from typing import Iterable, Iterator, List, Optional, Tuple
+from typing import Iterable, Iterator, List, Optional, Sequence, Tuple
 
 import numpy as np
 import torch
@@ -47,13 +49,12 @@ class SingleSceneDataset(DatasetBase, Configurable):
     def __len__(self) -> int:
         return len(self.poses)
 
-    # pyre-fixme[14]: `sequence_frames_in_order` overrides method defined in
-    #  `DatasetBase` inconsistently.
     def sequence_frames_in_order(
-        self, seq_name: str
+        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
     ) -> Iterator[Tuple[float, int, int]]:
         for i in range(len(self)):
-            yield (0.0, i, i)
+            if subset_filter is None or self.frame_types[i] in subset_filter:
+                yield 0.0, i, i
 
     def __getitem__(self, index) -> FrameData:
         if index >= len(self):

pytorch3d/implicitron/dataset/sql_dataset.py

@@ -0,0 +1,768 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import hashlib
+import json
+import logging
+import os
+from dataclasses import dataclass
+from typing import (
+    Any,
+    ClassVar,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+)
+
+import numpy as np
+import pandas as pd
+import sqlalchemy as sa
+import torch
+from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
+
+from pytorch3d.implicitron.dataset.frame_data import (  # noqa
+    FrameData,
+    FrameDataBuilder,
+    FrameDataBuilderBase,
+)
+from pytorch3d.implicitron.tools.config import (
+    registry,
+    ReplaceableBase,
+    run_auto_creation,
+)
+from sqlalchemy.orm import Session
+
+from .orm_types import SqlFrameAnnotation, SqlSequenceAnnotation
+
+
+logger = logging.getLogger(__name__)
+
+
+_SET_LISTS_TABLE: str = "set_lists"
+
+
+@registry.register
+class SqlIndexDataset(DatasetBase, ReplaceableBase):  # pyre-ignore
+    """
+    A dataset with annotations stored as SQLite tables. This is an index-based dataset.
+    The length is returned after all sequence and frame filters are applied (see param
+    definitions below). Indices can either be ordinal in [0, len), or pairs of
+    (sequence_name, frame_number); with the performance of `dataset[i]` and
+    `dataset[sequence_name, frame_number]` being same. A faster way to get metadata only
+    (without blobs) is `dataset.meta[idx]` indexing; it requires box_crop==False.
+    With ordinal indexing, the sequences are NOT guaranteed to span contiguous index
+    ranges, and frame numbers are NOT guaranteed to be increasing within a sequence.
+    Sequence-aware batch samplers have to use `sequence_[frames|indices]_in_order`
+    iterators, which are efficient.
+
+    This functionality requires SQLAlchemy 2.0 or later.
+
+    Metadata-related args:
+        sqlite_metadata_file: A SQLite file containing frame and sequence annotation
+            tables (mapping to SqlFrameAnnotation and SqlSequenceAnnotation,
+            respectively).
+        dataset_root: A root directory to look for images, masks, etc. It can be
+            alternatively set in `frame_data_builder` args, but this takes precedence.
+        subset_lists_file: A JSON/sqlite file containing the lists of frames
+            corresponding to different subsets (e.g. train/val/test) of the dataset;
+            format: {subset: [(sequence_name, frame_id, file_path)]}. All entries
+            must be present in frame_annotation metadata table.
+        path_manager: a facade for non-POSIX filesystems.
+        subsets: Restrict frames/sequences only to the given list of subsets
+            as defined in subset_lists_file (see above). Applied before all other
+            filters.
+        remove_empty_masks: Removes the frames with no active foreground pixels
+            in the segmentation mask (needs frame_annotation.mask.mass to be set;
+            null values are retained).
+        pick_frames_sql_clause: SQL WHERE clause to constrain frame annotations
+            NOTE: This is a potential security risk! The string is passed to the SQL
+            engine verbatim. Don’t expose it to end users of your application!
+        pick_categories: Restrict the dataset to the given list of categories.
+        pick_sequences: A Sequence of sequence names to restrict the dataset to.
+        exclude_sequences: A Sequence of the names of the sequences to exclude.
+        limit_sequences_per_category_to: Limit the dataset to the first up to N
+            sequences within each category (applies after all other sequence filters
+            but before `limit_sequences_to`).
+        limit_sequences_to: Limit the dataset to the first `limit_sequences_to`
+            sequences (after other sequence filters have been applied but before
+            frame-based filters).
+        limit_to: Limit the dataset to the first #limit_to frames (after other
+            filters have been applied, except n_frames_per_sequence).
+        n_frames_per_sequence: If > 0, randomly samples `n_frames_per_sequence`
+            frames in each sequences uniformly without replacement if it has
+            more frames than that; applied after other frame-level filters.
+        seed: The seed of the random generator sampling `n_frames_per_sequence`
+            random frames per sequence.
+    """
+
+    frame_annotations_type: ClassVar[Type[SqlFrameAnnotation]] = SqlFrameAnnotation
+
+    sqlite_metadata_file: str = ""
+    dataset_root: Optional[str] = None
+    subset_lists_file: str = ""
+    eval_batches_file: Optional[str] = None
+    path_manager: Any = None
+    subsets: Optional[List[str]] = None
+    remove_empty_masks: bool = True
+    pick_frames_sql_clause: Optional[str] = None
+    pick_categories: Tuple[str, ...] = ()
+
+    pick_sequences: Tuple[str, ...] = ()
+    exclude_sequences: Tuple[str, ...] = ()
+    limit_sequences_per_category_to: int = 0
+    limit_sequences_to: int = 0
+    limit_to: int = 0
+    n_frames_per_sequence: int = -1
+    seed: int = 0
+    remove_empty_masks_poll_whole_table_threshold: int = 300_000
+    # we set it manually in the constructor
+    # _index: pd.DataFrame = field(init=False)
+
+    frame_data_builder: FrameDataBuilderBase
+    frame_data_builder_class_type: str = "FrameDataBuilder"
+
+    def __post_init__(self) -> None:
+        if sa.__version__ < "2.0":
+            raise ImportError("This class requires SQL Alchemy 2.0 or later")
+
+        if not self.sqlite_metadata_file:
+            raise ValueError("sqlite_metadata_file must be set")
+
+        if self.dataset_root:
+            frame_builder_type = self.frame_data_builder_class_type
+            getattr(self, f"frame_data_builder_{frame_builder_type}_args")[
+                "dataset_root"
+            ] = self.dataset_root
+
+        run_auto_creation(self)
+        self.frame_data_builder.path_manager = self.path_manager
+
+        # pyre-ignore  # NOTE: sqlite-specific args (read-only mode).
+        self._sql_engine = sa.create_engine(
+            f"sqlite:///file:{self.sqlite_metadata_file}?mode=ro&uri=true"
+        )
+
+        sequences = self._get_filtered_sequences_if_any()
+
+        if self.subsets:
+            index = self._build_index_from_subset_lists(sequences)
+        else:
+            # TODO: if self.subset_lists_file and not self.subsets, it might be faster to
+            # still use the concatenated lists, assuming they cover the whole dataset
+            index = self._build_index_from_db(sequences)
+
+        if self.n_frames_per_sequence >= 0:
+            index = self._stratified_sample_index(index)
+
+        if len(index) == 0:
+            raise ValueError(f"There are no frames in the subsets: {self.subsets}!")
+
+        self._index = index.set_index(["sequence_name", "frame_number"])  # pyre-ignore
+
+        self.eval_batches = None  # pyre-ignore
+        if self.eval_batches_file:
+            self.eval_batches = self._load_filter_eval_batches()
+
+        logger.info(str(self))
+
+    def __len__(self) -> int:
+        # pyre-ignore[16]
+        return len(self._index)
+
+    def __getitem__(self, frame_idx: Union[int, Tuple[str, int]]) -> FrameData:
+        """
+        Fetches FrameData by either iloc in the index or by (sequence, frame_no) pair
+        """
+        return self._get_item(frame_idx, True)
+
+    @property
+    def meta(self):
+        """
+        Allows accessing metadata only without loading blobs using `dataset.meta[idx]`.
+        Requires box_crop==False, since in that case, cameras cannot be adjusted
+        without loading masks.
+
+        Returns:
+            FrameData objects with blob fields like `image_rgb` set to None.
+
+        Raises:
+            ValueError if dataset.box_crop is set.
+        """
+        return SqlIndexDataset._MetadataAccessor(self)
+
+    @dataclass
+    class _MetadataAccessor:
+        dataset: "SqlIndexDataset"
+
+        def __getitem__(self, frame_idx: Union[int, Tuple[str, int]]) -> FrameData:
+            return self.dataset._get_item(frame_idx, False)
+
+    def _get_item(
+        self, frame_idx: Union[int, Tuple[str, int]], load_blobs: bool = True
+    ) -> FrameData:
+        if isinstance(frame_idx, int):
+            if frame_idx >= len(self._index):
+                raise IndexError(f"index {frame_idx} out of range {len(self._index)}")
+
+            seq, frame = self._index.index[frame_idx]
+        else:
+            seq, frame, *rest = frame_idx
+            if isinstance(frame, torch.LongTensor):
+                frame = frame.item()
+
+            if (seq, frame) not in self._index.index:
+                raise IndexError(
+                    f"Sequence-frame index {frame_idx} not found; was it filtered out?"
+                )
+
+            if rest and rest[0] != self._index.loc[(seq, frame), "_image_path"]:
+                raise IndexError(f"Non-matching image path in {frame_idx}.")
+
+        stmt = sa.select(self.frame_annotations_type).where(
+            self.frame_annotations_type.sequence_name == seq,
+            self.frame_annotations_type.frame_number
+            == int(frame),  # cast from np.int64
+        )
+        seq_stmt = sa.select(SqlSequenceAnnotation).where(
+            SqlSequenceAnnotation.sequence_name == seq
+        )
+        with Session(self._sql_engine) as session:
+            entry = session.scalars(stmt).one()
+            seq_metadata = session.scalars(seq_stmt).one()
+
+        assert entry.image.path == self._index.loc[(seq, frame), "_image_path"]
+
+        frame_data = self.frame_data_builder.build(
+            entry, seq_metadata, load_blobs=load_blobs
+        )
+
+        # The rest of the fields are optional
+        frame_data.frame_type = self._get_frame_type(entry)
+        return frame_data
+
+    def __str__(self) -> str:
+        # pyre-ignore[16]
+        return f"SqlIndexDataset #frames={len(self._index)}"
+
+    def sequence_names(self) -> Iterable[str]:
+        """Returns an iterator over sequence names in the dataset."""
+        return self._index.index.unique("sequence_name")
+
+    # override
+    def category_to_sequence_names(self) -> Dict[str, List[str]]:
+        stmt = sa.select(
+            SqlSequenceAnnotation.category, SqlSequenceAnnotation.sequence_name
+        ).where(  # we limit results to sequences that have frames after all filters
+            SqlSequenceAnnotation.sequence_name.in_(self.sequence_names())
+        )
+        with self._sql_engine.connect() as connection:
+            cat_to_seqs = pd.read_sql(stmt, connection)
+
+        return cat_to_seqs.groupby("category")["sequence_name"].apply(list).to_dict()
+
+    # override
+    def get_frame_numbers_and_timestamps(
+        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
+    ) -> List[Tuple[int, float]]:
+        """
+        Implements the DatasetBase method.
+
+        NOTE: Avoid this function as there are more efficient alternatives such as
+        querying `dataset[idx]` directly or getting all sequence frames with
+        `sequence_[frames|indices]_in_order`.
+
+        Return the index and timestamp in their videos of the frames whose
+        indices are given in `idxs`. They need to belong to the same sequence!
+        If timestamps are absent, they are replaced with zeros.
+        This is used for letting SceneBatchSampler identify consecutive
+        frames.
+
+        Args:
+            idxs: a sequence int frame index in the dataset (it can be a slice)
+            subset_filter: must remain None
+
+        Returns:
+            list of tuples of
+                - frame index in video
+                - timestamp of frame in video, coalesced with 0s
+
+        Raises:
+            ValueError if idxs belong to more than one sequence.
+        """
+
+        if subset_filter is not None:
+            raise NotImplementedError(
+                "Subset filters are not supported in SQL Dataset. "
+                "We encourage creating a dataset per subset."
+            )
+
+        index_slice, _ = self._get_frame_no_coalesced_ts_by_row_indices(idxs)
+        # alternatively, we can use `.values.tolist()`, which may be faster
+        # but returns a list of lists
+        return list(index_slice.itertuples())
+
+    # override
+    def sequence_frames_in_order(
+        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
+    ) -> Iterator[Tuple[float, int, int]]:
+        """
+        Overrides the default DatasetBase implementation (we don’t use `_seq_to_idx`).
+        Returns an iterator over the frame indices in a given sequence.
+        We attempt to first sort by timestamp (if they are available),
+        then by frame number.
+
+        Args:
+            seq_name: the name of the sequence.
+            subset_filter: subset names to filter to
+
+        Returns:
+            an iterator over triplets `(timestamp, frame_no, dataset_idx)`,
+                where `frame_no` is the index within the sequence, and
+                `dataset_idx` is the index within the dataset.
+                `None` timestamps are replaced with 0s.
+        """
+        # TODO: implement sort_timestamp_first? (which would matter if the orders
+        # of frame numbers and timestamps are different)
+        rows = self._index.index.get_loc(seq_name)
+        if isinstance(rows, slice):
+            assert rows.stop is not None, "Unexpected result from pandas"
+            rows = range(rows.start or 0, rows.stop, rows.step or 1)
+        else:
+            rows = np.where(rows)[0]
+
+        index_slice, idx = self._get_frame_no_coalesced_ts_by_row_indices(
+            rows, seq_name, subset_filter
+        )
+        index_slice["idx"] = idx
+
+        yield from index_slice.itertuples(index=False)
+
+    # override
+    def get_eval_batches(self) -> Optional[List[Any]]:
+        """
+        This class does not support eval batches with ordinal indices. You can pass
+        eval_batches as a batch_sampler to a data_loader since the dataset supports
+        `dataset[seq_name, frame_no]` indexing.
+        """
+        return self.eval_batches
+
+    # override
+    def join(self, other_datasets: Iterable[DatasetBase]) -> None:
+        raise ValueError("Not supported! Preprocess the data by merging them instead.")
+
+    # override
+    @property
+    def frame_data_type(self) -> Type[FrameData]:
+        return self.frame_data_builder.frame_data_type
+
+    def is_filtered(self) -> bool:
+        """
+        Returns `True` in case the dataset has been filtered and thus some frame
+        annotations stored on the disk might be missing in the dataset object.
+        Does not account for subsets.
+
+        Returns:
+            is_filtered: `True` if the dataset has been filtered, else `False`.
+        """
+        return (
+            self.remove_empty_masks
+            or self.limit_to > 0
+            or self.limit_sequences_to > 0
+            or self.limit_sequences_per_category_to > 0
+            or len(self.pick_sequences) > 0
+            or len(self.exclude_sequences) > 0
+            or len(self.pick_categories) > 0
+            or self.n_frames_per_sequence > 0
+        )
+
+    def _get_filtered_sequences_if_any(self) -> Optional[pd.Series]:
+        # maximum possible filter (if limit_sequences_per_category_to == 0):
+        # WHERE category IN 'self.pick_categories'
+        # AND sequence_name IN 'self.pick_sequences'
+        # AND sequence_name NOT IN 'self.exclude_sequences'
+        # LIMIT 'self.limit_sequence_to'
+
+        where_conditions = [
+            *self._get_category_filters(),
+            *self._get_pick_filters(),
+            *self._get_exclude_filters(),
+        ]
+
+        def add_where(stmt):
+            return stmt.where(*where_conditions) if where_conditions else stmt
+
+        if self.limit_sequences_per_category_to <= 0:
+            stmt = add_where(sa.select(SqlSequenceAnnotation.sequence_name))
+        else:
+            subquery = sa.select(
+                SqlSequenceAnnotation.sequence_name,
+                sa.func.row_number()
+                .over(
+                    order_by=sa.text("ROWID"),  # NOTE: ROWID is SQLite-specific
+                    partition_by=SqlSequenceAnnotation.category,
+                )
+                .label("row_number"),
+            )
+
+            subquery = add_where(subquery).subquery()
+            stmt = sa.select(subquery.c.sequence_name).where(
+                subquery.c.row_number <= self.limit_sequences_per_category_to
+            )
+
+        if self.limit_sequences_to > 0:
+            logger.info(
+                f"Limiting dataset to first {self.limit_sequences_to} sequences"
+            )
+            # NOTE: ROWID is SQLite-specific
+            stmt = stmt.order_by(sa.text("ROWID")).limit(self.limit_sequences_to)
+
+        if (
+            not where_conditions
+            and self.limit_sequences_to <= 0
+            and self.limit_sequences_per_category_to <= 0
+        ):
+            # we will not need to filter by sequences
+            return None
+
+        with self._sql_engine.connect() as connection:
+            sequences = pd.read_sql_query(stmt, connection)["sequence_name"]
+        logger.info("... retained %d sequences" % len(sequences))
+
+        return sequences
+
+    def _get_category_filters(self) -> List[sa.ColumnElement]:
+        if not self.pick_categories:
+            return []
+
+        logger.info(f"Limiting dataset to categories: {self.pick_categories}")
+        return [SqlSequenceAnnotation.category.in_(self.pick_categories)]
+
+    def _get_pick_filters(self) -> List[sa.ColumnElement]:
+        if not self.pick_sequences:
+            return []
+
+        logger.info(f"Limiting dataset to sequences: {self.pick_sequences}")
+        return [SqlSequenceAnnotation.sequence_name.in_(self.pick_sequences)]
+
+    def _get_exclude_filters(self) -> List[sa.ColumnOperators]:
+        if not self.exclude_sequences:
+            return []
+
+        logger.info(f"Removing sequences from the dataset: {self.exclude_sequences}")
+        return [SqlSequenceAnnotation.sequence_name.notin_(self.exclude_sequences)]
+
+    def _load_subsets_from_json(self, subset_lists_path: str) -> pd.DataFrame:
+        assert self.subsets is not None
+        with open(subset_lists_path, "r") as f:
+            subset_to_seq_frame = json.load(f)
+
+        seq_frame_list = sum(
+            (
+                [(*row, subset) for row in subset_to_seq_frame[subset]]
+                for subset in self.subsets
+            ),
+            [],
+        )
+        index = pd.DataFrame(
+            seq_frame_list,
+            columns=["sequence_name", "frame_number", "_image_path", "subset"],
+        )
+        return index
+
+    def _load_subsets_from_sql(self, subset_lists_path: str) -> pd.DataFrame:
+        subsets = self.subsets
+        assert subsets is not None
+        # we need a new engine since we store the subsets in a separate DB
+        engine = sa.create_engine(f"sqlite:///{subset_lists_path}")
+        table = sa.Table(_SET_LISTS_TABLE, sa.MetaData(), autoload_with=engine)
+        stmt = sa.select(table).where(table.c.subset.in_(subsets))
+        with engine.connect() as connection:
+            index = pd.read_sql(stmt, connection)
+
+        return index
+
+    def _build_index_from_subset_lists(
+        self, sequences: Optional[pd.Series]
+    ) -> pd.DataFrame:
+        if not self.subset_lists_file:
+            raise ValueError("Requested subsets but subset_lists_file not given")
+
+        logger.info(f"Loading subset lists from {self.subset_lists_file}.")
+
+        subset_lists_path = self._local_path(self.subset_lists_file)
+        if subset_lists_path.lower().endswith(".json"):
+            index = self._load_subsets_from_json(subset_lists_path)
+        else:
+            index = self._load_subsets_from_sql(subset_lists_path)
+        index = index.set_index(["sequence_name", "frame_number"])
+        logger.info(f"  -> loaded {len(index)} samples of {self.subsets}.")
+
+        if sequences is not None:
+            logger.info("Applying filtered sequences.")
+            sequence_values = index.index.get_level_values("sequence_name")
+            index = index.loc[sequence_values.isin(sequences)]
+            logger.info(f"  -> retained {len(index)} samples.")
+
+        pick_frames_criteria = []
+        if self.remove_empty_masks:
+            logger.info("Culling samples with empty masks.")
+
+            if len(index) > self.remove_empty_masks_poll_whole_table_threshold:
+                # APPROACH 1: find empty masks and drop indices.
+                # dev load: 17s / 15 s (3.1M / 500K)
+                stmt = sa.select(
+                    self.frame_annotations_type.sequence_name,
+                    self.frame_annotations_type.frame_number,
+                ).where(self.frame_annotations_type._mask_mass == 0)
+                with Session(self._sql_engine) as session:
+                    to_remove = session.execute(stmt).all()
+
+                # Pandas uses np.int64 for integer types, so we have to case
+                # we might want to read it to pandas DataFrame directly to avoid the loop
+                to_remove = [(seq, np.int64(fr)) for seq, fr in to_remove]
+                index.drop(to_remove, errors="ignore", inplace=True)
+            else:
+                # APPROACH 3: load index into a temp table and join with annotations
+                # dev load: 94 s / 23 s (3.1M / 500K)
+                pick_frames_criteria.append(
+                    sa.or_(
+                        self.frame_annotations_type._mask_mass.is_(None),
+                        self.frame_annotations_type._mask_mass != 0,
+                    )
+                )
+
+        if self.pick_frames_sql_clause:
+            logger.info("Applying the custom SQL clause.")
+            pick_frames_criteria.append(sa.text(self.pick_frames_sql_clause))
+
+        if pick_frames_criteria:
+            index = self._pick_frames_by_criteria(index, pick_frames_criteria)
+
+        logger.info(f"  -> retained {len(index)} samples.")
+
+        if self.limit_to > 0:
+            logger.info(f"Limiting dataset to first {self.limit_to} frames")
+            index = index.sort_index().iloc[: self.limit_to]
+
+        return index.reset_index()
+
+    def _pick_frames_by_criteria(self, index: pd.DataFrame, criteria) -> pd.DataFrame:
+        IndexTable = self._get_temp_index_table_instance()
+        with self._sql_engine.connect() as connection:
+            IndexTable.create(connection)
+            # we don’t let pandas’s `to_sql` create the table automatically as
+            # the table would be permanent, so we create it and append with pandas
+            n_rows = index.to_sql(IndexTable.name, connection, if_exists="append")
+            assert n_rows == len(index)
+            sa_type = self.frame_annotations_type
+            stmt = (
+                sa.select(IndexTable)
+                .select_from(
+                    IndexTable.join(
+                        self.frame_annotations_type,
+                        sa.and_(
+                            sa_type.sequence_name == IndexTable.c.sequence_name,
+                            sa_type.frame_number == IndexTable.c.frame_number,
+                        ),
+                    )
+                )
+                .where(*criteria)
+            )
+            return pd.read_sql_query(stmt, connection).set_index(
+                ["sequence_name", "frame_number"]
+            )
+
+    def _build_index_from_db(self, sequences: Optional[pd.Series]):
+        logger.info("Loading sequcence-frame index from the database")
+        stmt = sa.select(
+            self.frame_annotations_type.sequence_name,
+            self.frame_annotations_type.frame_number,
+            self.frame_annotations_type._image_path,
+            sa.null().label("subset"),
+        )
+        where_conditions = []
+        if sequences is not None:
+            logger.info("  applying filtered sequences")
+            where_conditions.append(
+                self.frame_annotations_type.sequence_name.in_(sequences.tolist())
+            )
+
+        if self.remove_empty_masks:
+            logger.info("  excluding samples with empty masks")
+            where_conditions.append(
+                sa.or_(
+                    self.frame_annotations_type._mask_mass.is_(None),
+                    self.frame_annotations_type._mask_mass != 0,
+                )
+            )
+
+        if self.pick_frames_sql_clause:
+            logger.info("  applying custom SQL clause")
+            where_conditions.append(sa.text(self.pick_frames_sql_clause))
+
+        if where_conditions:
+            stmt = stmt.where(*where_conditions)
+
+        if self.limit_to > 0:
+            logger.info(f"Limiting dataset to first {self.limit_to} frames")
+            stmt = stmt.order_by(
+                self.frame_annotations_type.sequence_name,
+                self.frame_annotations_type.frame_number,
+            ).limit(self.limit_to)
+
+        with self._sql_engine.connect() as connection:
+            index = pd.read_sql_query(stmt, connection)
+
+        logger.info(f"  -> loaded {len(index)} samples.")
+        return index
+
+    def _sort_index_(self, index):
+        logger.info("Sorting the index by sequence and frame number.")
+        index.sort_values(["sequence_name", "frame_number"], inplace=True)
+        logger.info("  -> Done.")
+
+    def _load_filter_eval_batches(self):
+        assert self.eval_batches_file
+        logger.info(f"Loading eval batches from {self.eval_batches_file}")
+
+        if not os.path.isfile(self.eval_batches_file):
+            # The batch indices file does not exist.
+            # Most probably the user has not specified the root folder.
+            raise ValueError(
+                f"Looking for dataset json file in {self.eval_batches_file}. "
+                + "Please specify a correct dataset_root folder."
+            )
+
+        with open(self.eval_batches_file, "r") as f:
+            eval_batches = json.load(f)
+
+        # limit the dataset to sequences to allow multiple evaluations in one file
+        pick_sequences = set(self.pick_sequences)
+        if self.pick_categories:
+            cat_to_seq = self.category_to_sequence_names()
+            pick_sequences.update(
+                seq for cat in self.pick_categories for seq in cat_to_seq[cat]
+            )
+
+        if pick_sequences:
+            old_len = len(eval_batches)
+            eval_batches = [b for b in eval_batches if b[0][0] in pick_sequences]
+            logger.warn(
+                f"Picked eval batches by sequence/cat: {old_len} -> {len(eval_batches)}"
+            )
+
+        if self.exclude_sequences:
+            old_len = len(eval_batches)
+            exclude_sequences = set(self.exclude_sequences)
+            eval_batches = [b for b in eval_batches if b[0][0] not in exclude_sequences]
+            logger.warn(
+                f"Excluded eval batches by sequence: {old_len} -> {len(eval_batches)}"
+            )
+
+        return eval_batches
+
+    def _stratified_sample_index(self, index):
+        # NOTE this stratified sampling can be done more efficiently in
+        # the no-subset case above if it is added to the SQL query.
+        # We keep this generic implementation since no-subset case is uncommon
+        index = index.groupby("sequence_name", group_keys=False).apply(
+            lambda seq_frames: seq_frames.sample(
+                min(len(seq_frames), self.n_frames_per_sequence),
+                random_state=(
+                    _seq_name_to_seed(seq_frames.iloc[0]["sequence_name"]) + self.seed
+                ),
+            )
+        )
+        logger.info(f"  -> retained {len(index)} samples aster stratified sampling.")
+        return index
+
+    def _get_frame_type(self, entry: SqlFrameAnnotation) -> Optional[str]:
+        return self._index.loc[(entry.sequence_name, entry.frame_number), "subset"]
+
+    def _get_frame_no_coalesced_ts_by_row_indices(
+        self,
+        idxs: Sequence[int],
+        seq_name: Optional[str] = None,
+        subset_filter: Union[Sequence[str], str, None] = None,
+    ) -> Tuple[pd.DataFrame, Sequence[int]]:
+        """
+        Loads timestamps for given index rows belonging to the same sequence.
+        If seq_name is known, it speeds up the computation.
+        Raises ValueError if `idxs` do not all belong to a single sequences .
+        """
+        index_slice = self._index.iloc[idxs]
+        if subset_filter is not None:
+            if isinstance(subset_filter, str):
+                subset_filter = [subset_filter]
+            indicator = index_slice["subset"].isin(subset_filter)
+            index_slice = index_slice.loc[indicator]
+            idxs = [i for i, isin in zip(idxs, indicator) if isin]
+
+        frames = index_slice.index.get_level_values("frame_number").tolist()
+        if seq_name is None:
+            seq_name_list = index_slice.index.get_level_values("sequence_name").tolist()
+            seq_name_set = set(seq_name_list)
+            if len(seq_name_set) > 1:
+                raise ValueError("Given indices belong to more than one sequence.")
+            elif len(seq_name_set) == 1:
+                seq_name = seq_name_list[0]
+
+        coalesced_ts = sa.sql.functions.coalesce(
+            self.frame_annotations_type.frame_timestamp, 0
+        )
+        stmt = sa.select(
+            coalesced_ts.label("frame_timestamp"),
+            self.frame_annotations_type.frame_number,
+        ).where(
+            self.frame_annotations_type.sequence_name == seq_name,
+            self.frame_annotations_type.frame_number.in_(frames),
+        )
+
+        with self._sql_engine.connect() as connection:
+            frame_no_ts = pd.read_sql_query(stmt, connection)
+
+        if len(frame_no_ts) != len(index_slice):
+            raise ValueError(
+                "Not all indices are found in the database; "
+                "do they belong to more than one sequence?"
+            )
+
+        return frame_no_ts, idxs
+
+    def _local_path(self, path: str) -> str:
+        if self.path_manager is None:
+            return path
+        return self.path_manager.get_local_path(path)
+
+    def _get_temp_index_table_instance(self, table_name: str = "__index"):
+        CachedTable = self.frame_annotations_type.metadata.tables.get(table_name)
+        if CachedTable is not None:  # table definition is not idempotent
+            return CachedTable
+
+        return sa.Table(
+            table_name,
+            self.frame_annotations_type.metadata,
+            sa.Column("sequence_name", sa.String, primary_key=True),
+            sa.Column("frame_number", sa.Integer, primary_key=True),
+            sa.Column("_image_path", sa.String),
+            sa.Column("subset", sa.String),
+            prefixes=["TEMP"],  # NOTE SQLite specific!
+        )
+
+
+def _seq_name_to_seed(seq_name) -> int:
+    """Generates numbers in [0, 2 ** 28)"""
+    return int(hashlib.sha1(seq_name.encode("utf-8")).hexdigest()[:7], 16)
+
+
+def _safe_as_tensor(data, dtype):
+    return torch.tensor(data, dtype=dtype) if data is not None else None

pytorch3d/implicitron/dataset/sql_dataset_provider.py

@@ -0,0 +1,424 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import logging
+import os
+from typing import List, Optional, Tuple, Type
+
+import numpy as np
+
+from omegaconf import DictConfig, OmegaConf
+
+from pytorch3d.implicitron.dataset.dataset_map_provider import (
+    DatasetMap,
+    DatasetMapProviderBase,
+    PathManagerFactory,
+)
+from pytorch3d.implicitron.tools.config import (
+    expand_args_fields,
+    registry,
+    run_auto_creation,
+)
+
+from .sql_dataset import SqlIndexDataset
+
+
+_CO3D_SQL_DATASET_ROOT: str = os.getenv("CO3D_SQL_DATASET_ROOT", "")
+
+# _NEED_CONTROL is a list of those elements of SqlIndexDataset which
+# are not directly specified for it in the config but come from the
+# DatasetMapProvider.
+_NEED_CONTROL: Tuple[str, ...] = (
+    "path_manager",
+    "subsets",
+    "sqlite_metadata_file",
+    "subset_lists_file",
+)
+
+logger = logging.getLogger(__name__)
+
+
+@registry.register
+class SqlIndexDatasetMapProvider(DatasetMapProviderBase):  # pyre-ignore [13]
+    """
+    Generates the training, validation, and testing dataset objects for
+    a dataset laid out on disk like SQL-CO3D, with annotations in an SQLite data base.
+
+    The dataset is organized in the filesystem as follows::
+
+        self.dataset_root
+            ├── <possible/partition/0>
+            │   ├── <sequence_name_0>
+            │   │   ├── depth_masks
+            │   │   ├── depths
+            │   │   ├── images
+            │   │   ├── masks
+            │   │   └── pointcloud.ply
+            │   ├── <sequence_name_1>
+            │   │   ├── depth_masks
+            │   │   ├── depths
+            │   │   ├── images
+            │   │   ├── masks
+            │   │   └── pointcloud.ply
+            │   ├── ...
+            │   ├── <sequence_name_N>
+            │   ├── set_lists
+            │       ├── <subset_base_name_0>.json
+            │       ├── <subset_base_name_1>.json
+            │       ├── ...
+            │       ├── <subset_base_name_2>.json
+            │   ├── eval_batches
+            │   │   ├── <eval_batches_base_name_0>.json
+            │   │   ├── <eval_batches_base_name_1>.json
+            │   │   ├── ...
+            │   │   ├── <eval_batches_base_name_M>.json
+            │   ├── frame_annotations.jgz
+            │   ├── sequence_annotations.jgz
+            ├── <possible/partition/1>
+            ├── ...
+            ├── <possible/partition/K>
+            ├── set_lists
+                ├── <subset_base_name_0>.sqlite
+                ├── <subset_base_name_1>.sqlite
+                ├── ...
+                ├── <subset_base_name_2>.sqlite
+            ├── eval_batches
+            │   ├── <eval_batches_base_name_0>.json
+            │   ├── <eval_batches_base_name_1>.json
+            │   ├── ...
+            │   ├── <eval_batches_base_name_M>.json
+
+    The dataset contains sequences named `<sequence_name_i>` that may be partitioned by
+    directories such as `<possible/partition/0>` e.g. representing categories but they
+    can also be stored in a flat structure. Each sequence folder contains the list of
+    sequence images, depth maps, foreground masks, and valid-depth masks
+    `images`, `depths`, `masks`, and `depth_masks` respectively. Furthermore,
+    `set_lists/` dirtectories (with partitions or global) store json or sqlite files
+    `<subset_base_name_l>.<ext>`, each describing a certain sequence subset.
+    These subset path conventions are not hard-coded and arbitrary relative path can be
+    specified by setting `self.subset_lists_path` to the relative path w.r.t.
+    dataset root.
+
+    Each `<subset_base_name_l>.json` file contains the following dictionary::
+
+        {
+            "train": [
+                (sequence_name: str, frame_number: int, image_path: str),
+                ...
+            ],
+            "val": [
+                (sequence_name: str, frame_number: int, image_path: str),
+                ...
+            ],
+            "test": [
+                (sequence_name: str, frame_number: int, image_path: str),
+                ...
+            ],
+        ]
+
+    defining the list of frames (identified with their `sequence_name` and
+    `frame_number`) in the "train", "val", and "test" subsets of the dataset. In case of
+    SQLite format, `<subset_base_name_l>.sqlite` contains a table with the header::
+
+        | sequence_name | frame_number | image_path | subset |
+
+    Note that `frame_number` can be obtained only from the metadata and
+    does not necesarrily correspond to the numeric suffix of the corresponding image
+    file name (e.g. a file `<partition_0>/<sequence_name_0>/images/frame00005.jpg` can
+    have its frame number set to `20`, not 5).
+
+    Each `<eval_batches_base_name_M>.json` file contains a list of evaluation examples
+    in the following form::
+
+        [
+            [  # batch 1
+                (sequence_name: str, frame_number: int, image_path: str),
+                ...
+            ],
+            [  # batch 2
+                (sequence_name: str, frame_number: int, image_path: str),
+                ...
+            ],
+        ]
+
+    Note that the evaluation examples always come from the `"test"` subset of the dataset.
+    (test frames can repeat across batches). The batches can contain single element,
+    which is typical in case of regular radiance field fitting.
+
+    Args:
+        subset_lists_path: The relative path to the dataset subset definition.
+            For CO3D, these include e.g. "skateboard/set_lists/set_lists_manyview_dev_0.json".
+            By default (None), dataset is not partitioned to subsets (in that case, setting
+            `ignore_subsets` will speed up construction)
+        dataset_root: The root folder of the dataset.
+        metadata_basename: name of the SQL metadata file in dataset_root;
+            not expected to be changed by users
+        test_on_train: Construct validation and test datasets from
+            the training subset; note that in practice, in this
+            case all subset dataset objects will be same
+        only_test_set: Load only the test set. Incompatible with `test_on_train`.
+        ignore_subsets: Don’t filter by subsets in the dataset; note that in this
+            case all subset datasets will be same
+        eval_batch_num_training_frames: Add a certain number of training frames to each
+            eval batch. Useful for evaluating models that require
+            source views as input (e.g. NeRF-WCE / PixelNeRF).
+        dataset_args: Specifies additional arguments to the
+            JsonIndexDataset constructor call.
+        path_manager_factory: (Optional) An object that generates an instance of
+            PathManager that can translate provided file paths.
+        path_manager_factory_class_type: The class type of `path_manager_factory`.
+    """
+
+    category: Optional[str] = None
+    subset_list_name: Optional[str] = None  # TODO: docs
+    # OR
+    subset_lists_path: Optional[str] = None
+    eval_batches_path: Optional[str] = None
+
+    dataset_root: str = _CO3D_SQL_DATASET_ROOT
+    metadata_basename: str = "metadata.sqlite"
+
+    test_on_train: bool = False
+    only_test_set: bool = False
+    ignore_subsets: bool = False
+    train_subsets: Tuple[str, ...] = ("train",)
+    val_subsets: Tuple[str, ...] = ("val",)
+    test_subsets: Tuple[str, ...] = ("test",)
+
+    eval_batch_num_training_frames: int = 0
+
+    # this is a mould that is never constructed, used to build self._dataset_map values
+    dataset_class_type: str = "SqlIndexDataset"
+    dataset: SqlIndexDataset
+
+    path_manager_factory: PathManagerFactory
+    path_manager_factory_class_type: str = "PathManagerFactory"
+
+    def __post_init__(self):
+        super().__init__()
+        run_auto_creation(self)
+
+        if self.only_test_set and self.test_on_train:
+            raise ValueError("Cannot have only_test_set and test_on_train")
+
+        if self.ignore_subsets and not self.only_test_set:
+            self.test_on_train = True  # no point in loading same data 3 times
+
+        path_manager = self.path_manager_factory.get()
+
+        sqlite_metadata_file = os.path.join(self.dataset_root, self.metadata_basename)
+        sqlite_metadata_file = _local_path(path_manager, sqlite_metadata_file)
+
+        if not os.path.isfile(sqlite_metadata_file):
+            # The sqlite_metadata_file does not exist.
+            # Most probably the user has not specified the root folder.
+            raise ValueError(
+                f"Looking for frame annotations in {sqlite_metadata_file}."
+                + " Please specify a correct dataset_root folder."
+                + " Note: By default the root folder is taken from the"
+                + " CO3D_SQL_DATASET_ROOT environment variable."
+            )
+
+        if self.subset_lists_path and self.subset_list_name:
+            raise ValueError(
+                "subset_lists_path and subset_list_name cannot be both set"
+            )
+
+        subset_lists_file = self._get_lists_file("set_lists")
+
+        # setup the common dataset arguments
+        common_dataset_kwargs = {
+            **getattr(self, f"dataset_{self.dataset_class_type}_args"),
+            "sqlite_metadata_file": sqlite_metadata_file,
+            "dataset_root": self.dataset_root,
+            "subset_lists_file": subset_lists_file,
+            "path_manager": path_manager,
+        }
+
+        if self.category:
+            logger.info(f"Forcing category filter in the datasets to {self.category}")
+            common_dataset_kwargs["pick_categories"] = self.category.split(",")
+
+        # get the used dataset type
+        dataset_type: Type[SqlIndexDataset] = registry.get(
+            SqlIndexDataset, self.dataset_class_type
+        )
+        expand_args_fields(dataset_type)
+
+        if subset_lists_file is not None and not os.path.isfile(subset_lists_file):
+            available_subsets = self._get_available_subsets(
+                OmegaConf.to_object(common_dataset_kwargs["pick_categories"])
+            )
+            msg = f"Cannot find subset list file {self.subset_lists_path}."
+            if available_subsets:
+                msg += f" Some of the available subsets: {str(available_subsets)}."
+            raise ValueError(msg)
+
+        train_dataset = None
+        val_dataset = None
+        if not self.only_test_set:
+            # load the training set
+            logger.debug("Constructing train dataset.")
+            train_dataset = dataset_type(
+                **common_dataset_kwargs, subsets=self._get_subsets(self.train_subsets)
+            )
+            logger.info(f"Train dataset: {str(train_dataset)}")
+
+        if self.test_on_train:
+            assert train_dataset is not None
+            val_dataset = test_dataset = train_dataset
+        else:
+            # load the val and test sets
+            if not self.only_test_set:
+                # NOTE: this is always loaded in JsonProviderV2
+                logger.debug("Extracting val dataset.")
+                val_dataset = dataset_type(
+                    **common_dataset_kwargs, subsets=self._get_subsets(self.val_subsets)
+                )
+                logger.info(f"Val dataset: {str(val_dataset)}")
+
+            logger.debug("Extracting test dataset.")
+            eval_batches_file = self._get_lists_file("eval_batches")
+            del common_dataset_kwargs["eval_batches_file"]
+            test_dataset = dataset_type(
+                **common_dataset_kwargs,
+                subsets=self._get_subsets(self.test_subsets, True),
+                eval_batches_file=eval_batches_file,
+            )
+            logger.info(f"Test dataset: {str(test_dataset)}")
+
+            if (
+                eval_batches_file is not None
+                and self.eval_batch_num_training_frames > 0
+            ):
+                self._extend_eval_batches(test_dataset)
+
+        self._dataset_map = DatasetMap(
+            train=train_dataset, val=val_dataset, test=test_dataset
+        )
+
+    def _get_subsets(self, subsets, is_eval: bool = False):
+        if self.ignore_subsets:
+            return None
+
+        if is_eval and self.eval_batch_num_training_frames > 0:
+            # we will need to have training frames for extended batches
+            return list(subsets) + list(self.train_subsets)
+
+        return subsets
+
+    def _extend_eval_batches(self, test_dataset: SqlIndexDataset) -> None:
+        rng = np.random.default_rng(seed=0)
+        eval_batches = test_dataset.get_eval_batches()
+        if eval_batches is None:
+            raise ValueError("Eval batches were not loaded!")
+
+        for batch in eval_batches:
+            sequence = batch[0][0]
+            seq_frames = list(
+                test_dataset.sequence_frames_in_order(sequence, self.train_subsets)
+            )
+            idx_to_add = rng.permutation(len(seq_frames))[
+                : self.eval_batch_num_training_frames
+            ]
+            batch.extend((sequence, seq_frames[a][1]) for a in idx_to_add)
+
+    @classmethod
+    def dataset_tweak_args(cls, type, args: DictConfig) -> None:
+        """
+        Called by get_default_args.
+        Certain fields are not exposed on each dataset class
+        but rather are controlled by this provider class.
+        """
+        for key in _NEED_CONTROL:
+            del args[key]
+
+    def create_dataset(self):
+        # No `dataset` member of this class is created.
+        # The dataset(s) live in `self.get_dataset_map`.
+        pass
+
+    def get_dataset_map(self) -> DatasetMap:
+        return self._dataset_map  # pyre-ignore [16]
+
+    def _get_available_subsets(self, categories: List[str]):
+        """
+        Get the available subset names for a given category folder (if given) inside
+        a root dataset folder `dataset_root`.
+        """
+        path_manager = self.path_manager_factory.get()
+
+        subsets: List[str] = []
+        for prefix in [""] + categories:
+            set_list_dir = os.path.join(self.dataset_root, prefix, "set_lists")
+            if not (
+                (path_manager is not None) and path_manager.isdir(set_list_dir)
+            ) and not os.path.isdir(set_list_dir):
+                continue
+
+            set_list_files = (os.listdir if path_manager is None else path_manager.ls)(
+                set_list_dir
+            )
+            subsets.extend(os.path.join(prefix, "set_lists", f) for f in set_list_files)
+
+        return subsets
+
+    def _get_lists_file(self, flavor: str) -> Optional[str]:
+        if flavor == "eval_batches":
+            subset_lists_path = self.eval_batches_path
+        else:
+            subset_lists_path = self.subset_lists_path
+
+        if not subset_lists_path and not self.subset_list_name:
+            return None
+
+        category_elem = ""
+        if self.category and "," not in self.category:
+            # if multiple categories are given, looking for global set lists
+            category_elem = self.category
+
+        subset_lists_path = subset_lists_path or (
+            os.path.join(
+                category_elem, f"{flavor}", f"{flavor}_{self.subset_list_name}"
+            )
+        )
+
+        assert subset_lists_path
+        path_manager = self.path_manager_factory.get()
+        # try absolute path first
+        subset_lists_file = _get_local_path_check_extensions(
+            subset_lists_path, path_manager
+        )
+        if subset_lists_file:
+            return subset_lists_file
+
+        full_path = os.path.join(self.dataset_root, subset_lists_path)
+        subset_lists_file = _get_local_path_check_extensions(full_path, path_manager)
+
+        if not subset_lists_file:
+            raise FileNotFoundError(
+                f"Subset lists path given but not found: {full_path}"
+            )
+
+        return subset_lists_file
+
+
+def _get_local_path_check_extensions(
+    path, path_manager, extensions=("", ".sqlite", ".json")
+) -> Optional[str]:
+    for ext in extensions:
+        local = _local_path(path_manager, path + ext)
+        if os.path.isfile(local):
+            return local
+
+    return None
+
+
+def _local_path(path_manager, path: str) -> str:
+    if path_manager is None:
+        return path
+    return path_manager.get_local_path(path)

pytorch3d/implicitron/dataset/train_eval_data_loader_provider.py

@@ -0,0 +1,191 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import logging
+from typing import Any, Dict, Optional, Tuple
+
+from pytorch3d.implicitron.dataset.data_loader_map_provider import (
+    DataLoaderMap,
+    SceneBatchSampler,
+    SequenceDataLoaderMapProvider,
+)
+from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
+from pytorch3d.implicitron.dataset.dataset_map_provider import DatasetMap
+from pytorch3d.implicitron.dataset.frame_data import FrameData
+from pytorch3d.implicitron.tools.config import registry, run_auto_creation
+
+from torch.utils.data import DataLoader
+
+logger = logging.getLogger(__name__)
+
+
+# TODO: we can merge it with SequenceDataLoaderMapProvider in PyTorch3D
+# and support both eval_batches protocols
+@registry.register
+class TrainEvalDataLoaderMapProvider(SequenceDataLoaderMapProvider):
+    """
+    Implementation of DataLoaderMapProviderBase that may use internal eval batches for
+    the test dataset. In particular, if `eval_batches_relpath` is set, it loads
+    eval batches from that json file, otherwise test set is treated in the same way as
+    train and val, i.e. the parameters `dataset_length_test` and `test_conditioning_type`
+    are respected.
+
+    If conditioning is not required, then the batch size should
+    be set as 1, and most of the fields do not matter.
+
+    If conditioning is required, each batch will contain one main
+    frame first to predict and the, rest of the elements are for
+    conditioning.
+
+    If images_per_seq_options is left empty, the conditioning
+    frames are picked according to the conditioning type given.
+    This does not have regard to the order of frames in a
+    scene, or which frames belong to what scene.
+
+    If images_per_seq_options is given, then the conditioning types
+    must be SAME and the remaining fields are used.
+
+    Members:
+        batch_size: The size of the batch of the data loader.
+        num_workers: Number of data-loading threads in each data loader.
+        dataset_length_train: The number of batches in a training epoch. Or 0 to mean
+            an epoch is the length of the training set.
+        dataset_length_val: The number of batches in a validation epoch. Or 0 to mean
+            an epoch is the length of the validation set.
+        dataset_length_test: used if test_dataset.eval_batches is NOT set. The number of
+            batches in a testing epoch. Or 0 to mean an epoch is the length of the test
+            set.
+        images_per_seq_options: Possible numbers of frames sampled per sequence in a batch.
+            If a conditioning_type is KNOWN or TRAIN, then this must be left at its initial
+            value. Empty (the default) means that we are not careful about which frames
+            come from which scene.
+        sample_consecutive_frames: if True, will sample a contiguous interval of frames
+            in the sequence. It first sorts the frames by timestimps when available,
+            otherwise by frame numbers, finds the connected segments within the sequence
+            of sufficient length, then samples a random pivot element among them and
+            ideally uses it as a middle of the temporal window, shifting the borders
+            where necessary. This strategy mitigates the bias against shorter segments
+            and their boundaries.
+        consecutive_frames_max_gap: if a number > 0, then used to define the maximum
+            difference in frame_number of neighbouring frames when forming connected
+            segments; if both this and consecutive_frames_max_gap_seconds are 0s,
+            the whole sequence is considered a segment regardless of frame numbers.
+        consecutive_frames_max_gap_seconds: if a number > 0.0, then used to define the
+            maximum difference in frame_timestamp of neighbouring frames when forming
+            connected segments; if both this and consecutive_frames_max_gap are 0s,
+            the whole sequence is considered a segment regardless of frame timestamps.
+    """
+
+    batch_size: int = 1
+    num_workers: int = 0
+
+    dataset_length_train: int = 0
+    dataset_length_val: int = 0
+    dataset_length_test: int = 0
+
+    images_per_seq_options: Tuple[int, ...] = ()
+    sample_consecutive_frames: bool = False
+    consecutive_frames_max_gap: int = 0
+    consecutive_frames_max_gap_seconds: float = 0.1
+
+    def __post_init__(self):
+        run_auto_creation(self)
+
+    def get_data_loader_map(self, datasets: DatasetMap) -> DataLoaderMap:
+        """
+        Returns a collection of data loaders for a given collection of datasets.
+        """
+        train = self._make_generic_data_loader(
+            datasets.train,
+            self.dataset_length_train,
+            datasets.train,
+        )
+
+        val = self._make_generic_data_loader(
+            datasets.val,
+            self.dataset_length_val,
+            datasets.train,
+        )
+
+        if datasets.test is not None and datasets.test.get_eval_batches() is not None:
+            test = self._make_eval_data_loader(datasets.test)
+        else:
+            test = self._make_generic_data_loader(
+                datasets.test,
+                self.dataset_length_test,
+                datasets.train,
+            )
+
+        return DataLoaderMap(train=train, val=val, test=test)
+
+    def _make_eval_data_loader(
+        self,
+        dataset: Optional[DatasetBase],
+    ) -> Optional[DataLoader[FrameData]]:
+        if dataset is None:
+            return None
+
+        return DataLoader(
+            dataset,
+            batch_sampler=dataset.get_eval_batches(),
+            **self._get_data_loader_common_kwargs(dataset),
+        )
+
+    def _make_generic_data_loader(
+        self,
+        dataset: Optional[DatasetBase],
+        num_batches: int,
+        train_dataset: Optional[DatasetBase],
+    ) -> Optional[DataLoader[FrameData]]:
+        """
+        Returns the dataloader for a dataset.
+
+        Args:
+            dataset: the dataset
+            num_batches: possible ceiling on number of batches per epoch
+            train_dataset: the training dataset, used if conditioning_type==TRAIN
+            conditioning_type: source for padding of batches
+        """
+        if dataset is None:
+            return None
+
+        data_loader_kwargs = self._get_data_loader_common_kwargs(dataset)
+
+        if len(self.images_per_seq_options) > 0:
+            # this is a typical few-view setup
+            # conditioning comes from the same subset since subsets are split by seqs
+            batch_sampler = SceneBatchSampler(
+                dataset,
+                self.batch_size,
+                num_batches=len(dataset) if num_batches <= 0 else num_batches,
+                images_per_seq_options=self.images_per_seq_options,
+                sample_consecutive_frames=self.sample_consecutive_frames,
+                consecutive_frames_max_gap=self.consecutive_frames_max_gap,
+                consecutive_frames_max_gap_seconds=self.consecutive_frames_max_gap_seconds,
+            )
+            return DataLoader(
+                dataset,
+                batch_sampler=batch_sampler,
+                **data_loader_kwargs,
+            )
+
+        if self.batch_size == 1:
+            # this is a typical many-view setup (without conditioning)
+            return self._simple_loader(dataset, num_batches, data_loader_kwargs)
+
+        # edge case: conditioning on train subset, typical for Nerformer-like many-view
+        # there is only one sequence in all datasets, so we condition on another subset
+        return self._train_loader(
+            dataset, train_dataset, num_batches, data_loader_kwargs
+        )
+
+    def _get_data_loader_common_kwargs(self, dataset: DatasetBase) -> Dict[str, Any]:
+        return {
+            "num_workers": self.num_workers,
+            "collate_fn": dataset.frame_data_type.collate,
+        }