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,28 +156,54 @@ 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:
+          conda_docker_image: pytorch/conda-builder:cuda113
+          context: DOCKERHUB_TOKEN
+          cu_version: cu113
+          name: linux_conda_py38_cu113_pyt1120
+          python_version: '3.8'
+          pytorch_version: 1.12.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py38_cu116_pyt1120
+          python_version: '3.8'
+          pytorch_version: 1.12.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda113
+          context: DOCKERHUB_TOKEN
+          cu_version: cu113
+          name: linux_conda_py38_cu113_pyt1121
+          python_version: '3.8'
+          pytorch_version: 1.12.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py38_cu116_pyt1121
+          python_version: '3.8'
+          pytorch_version: 1.12.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py38_cu116_pyt1130
+          python_version: '3.8'
+          pytorch_version: 1.13.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py38_cu117_pyt1130
+          python_version: '3.8'
+          pytorch_version: 1.13.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda116
           context: DOCKERHUB_TOKEN
@@ -192,6 +218,160 @@ workflows:
           name: linux_conda_py38_cu117_pyt1131
           python_version: '3.8'
           pytorch_version: 1.13.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py38_cu117_pyt200
+          python_version: '3.8'
+          pytorch_version: 2.0.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py38_cu118_pyt200
+          python_version: '3.8'
+          pytorch_version: 2.0.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          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: 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: 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: 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: 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: 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:cuda118
+          context: DOCKERHUB_TOKEN
+          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: 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: 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:cuda121
+          context: DOCKERHUB_TOKEN
+          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: 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: 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:cuda118
+          context: DOCKERHUB_TOKEN
+          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: 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
+          cu_version: cu113
+          name: linux_conda_py39_cu113_pyt1120
+          python_version: '3.9'
+          pytorch_version: 1.12.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py39_cu116_pyt1120
+          python_version: '3.9'
+          pytorch_version: 1.12.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda113
+          context: DOCKERHUB_TOKEN
+          cu_version: cu113
+          name: linux_conda_py39_cu113_pyt1121
+          python_version: '3.9'
+          pytorch_version: 1.12.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py39_cu116_pyt1121
+          python_version: '3.9'
+          pytorch_version: 1.12.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py39_cu116_pyt1130
+          python_version: '3.9'
+          pytorch_version: 1.13.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py39_cu117_pyt1130
+          python_version: '3.9'
+          pytorch_version: 1.13.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda116
           context: DOCKERHUB_TOKEN
@@ -206,6 +386,160 @@ workflows:
           name: linux_conda_py39_cu117_pyt1131
           python_version: '3.9'
           pytorch_version: 1.13.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py39_cu117_pyt200
+          python_version: '3.9'
+          pytorch_version: 2.0.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py39_cu118_pyt200
+          python_version: '3.9'
+          pytorch_version: 2.0.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          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: 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:cuda118
+          context: DOCKERHUB_TOKEN
+          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:cuda121
+          context: DOCKERHUB_TOKEN
+          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: 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
+          cu_version: cu113
+          name: linux_conda_py310_cu113_pyt1120
+          python_version: '3.10'
+          pytorch_version: 1.12.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py310_cu116_pyt1120
+          python_version: '3.10'
+          pytorch_version: 1.12.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda113
+          context: DOCKERHUB_TOKEN
+          cu_version: cu113
+          name: linux_conda_py310_cu113_pyt1121
+          python_version: '3.10'
+          pytorch_version: 1.12.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py310_cu116_pyt1121
+          python_version: '3.10'
+          pytorch_version: 1.12.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda116
+          context: DOCKERHUB_TOKEN
+          cu_version: cu116
+          name: linux_conda_py310_cu116_pyt1130
+          python_version: '3.10'
+          pytorch_version: 1.13.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py310_cu117_pyt1130
+          python_version: '3.10'
+          pytorch_version: 1.13.0
       - binary_linux_conda:
           conda_docker_image: pytorch/conda-builder:cuda116
           context: DOCKERHUB_TOKEN
@@ -220,3 +554,247 @@ workflows:
           name: linux_conda_py310_cu117_pyt1131
           python_version: '3.10'
           pytorch_version: 1.13.1
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py310_cu117_pyt200
+          python_version: '3.10'
+          pytorch_version: 2.0.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda118
+          context: DOCKERHUB_TOKEN
+          cu_version: cu118
+          name: linux_conda_py310_cu118_pyt200
+          python_version: '3.10'
+          pytorch_version: 2.0.0
+      - binary_linux_conda:
+          conda_docker_image: pytorch/conda-builder:cuda117
+          context: DOCKERHUB_TOKEN
+          cu_version: cu117
+          name: linux_conda_py310_cu117_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_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,55 +18,55 @@ 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 cuda_version == "cu113":
-        return "pytorch/conda-builder:cuda113"
-    if cuda_version == "cu115":
-        return "pytorch/conda-builder:cuda115"
-    if cuda_version == "cu116":
-        return "pytorch/conda-builder:cuda116"
-    if cuda_version == "cu117":
-        return "pytorch/conda-builder:cuda117"
-    raise ValueError("Unknown cuda version")
+    if len(cuda_version) != 5:
+        raise ValueError("Unknown cuda version")
+    return "pytorch/conda-builder:cuda" + cuda_version[2:]
 
 
 def pytorch_versions_for_python(python_version):
-    if python_version in ["3.7", "3.8"]:
+    if python_version in ["3.8", "3.9"]:
         return list(CONDA_CUDA_VERSIONS)
-    if python_version == "3.9":
-        return [
-            i
-            for i in CONDA_CUDA_VERSIONS
-            if version.Version(i) > version.Version("1.7.0")
-        ]
     if python_version == "3.10":
         return [
             i
             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,19 +9,19 @@ 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 or 1.13.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)
 - [ioPath](https://github.com/facebookresearch/iopath)
 - If CUDA is to be used, use a version which is supported by the corresponding pytorch version and at least version 9.2.
-- If CUDA is to be used and you are building from source, the CUB library must be available. We recommend version 1.10.0.
+- If CUDA older than 11.7 is to be used and you are building from source, the CUB library must be available. We recommend version 1.10.0.
 
 The runtime dependencies can be installed by running:
 ```
 conda create -n pytorch3d python=3.9
 conda activate pytorch3d
-conda install -c pytorch pytorch=1.9.1 torchvision cudatoolkit=11.6
+conda install pytorch=1.13.0 torchvision pytorch-cuda=11.6 -c pytorch -c nvidia
 conda install -c fvcore -c iopath -c conda-forge fvcore iopath
 ```
 
@@ -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.
@@ -102,6 +97,7 @@ version_str="".join([
     torch.version.cuda.replace(".",""),
     f"_pyt{pyt_version_str}"
 ])
+!pip install fvcore iopath
 !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
 ```
 

README.md

@@ -12,7 +12,7 @@ Key features include:
 - Data structure for storing and manipulating triangle meshes
 - Efficient operations on triangle meshes (projective transformations, graph convolution, sampling, loss functions)
 - A differentiable mesh renderer
-- Implicitron, see [its README](projects/implicitron_trainer), a framework for new-view synthesis via implicit representations.
+- Implicitron, see [its README](projects/implicitron_trainer), a framework for new-view synthesis via implicit representations. ([blog post](https://ai.facebook.com/blog/implicitron-a-new-modular-extensible-framework-for-neural-implicit-representations-in-pytorch3d/))
 
 PyTorch3D is designed to integrate smoothly with deep learning methods for predicting and manipulating 3D data.
 For this reason, all operators in PyTorch3D:
@@ -24,6 +24,8 @@ For this reason, all operators in PyTorch3D:
 
 Within FAIR, PyTorch3D has been used to power research projects such as [Mesh R-CNN](https://arxiv.org/abs/1906.02739).
 
+See our [blog post](https://ai.facebook.com/blog/-introducing-pytorch3d-an-open-source-library-for-3d-deep-learning/) to see more demos and learn about PyTorch3D.
+
 ## Installation
 
 For detailed instructions refer to [INSTALL.md](INSTALL.md).
@@ -144,6 +146,14 @@ 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.
 
 **[Aug 10th 2022]:**   PyTorch3D [v0.7.0](https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.0) released with Implicitron and MeshRasterizerOpenGL.

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/notes/io.md

@@ -7,20 +7,20 @@ sidebar_label: File IO
 There is a flexible interface for loading and saving point clouds and meshes from different formats.
 
 The main usage is via the `pytorch3d.io.IO` object, and its methods
-`load_mesh`, `save_mesh`, `load_point_cloud` and `save_point_cloud`.
+`load_mesh`, `save_mesh`, `load_pointcloud` and `save_pointcloud`.
 
 For example, to load a mesh you might do
 ```
 from pytorch3d.io import IO
 
 device=torch.device("cuda:0")
-mesh = IO().load_mesh("mymesh.ply", device=device)
+mesh = IO().load_mesh("mymesh.obj", device=device)
 ```
 
 and to save a pointcloud you might do
 ```
 pcl = Pointclouds(...)
-IO().save_point_cloud(pcl, "output_pointcloud.obj")
+IO().save_pointcloud(pcl, "output_pointcloud.ply")
 ```
 
 For meshes, this supports OBJ, PLY and OFF files.
@@ -31,4 +31,4 @@ In addition, there is experimental support for loading meshes from
 [glTF 2 assets](https://github.com/KhronosGroup/glTF/tree/master/specification/2.0)
 stored either in a GLB container file or a glTF JSON file with embedded binary data.
 This must be enabled explicitly, as described in
-`pytorch3d/io/experimental_gltf_io.ply`.
+`pytorch3d/io/experimental_gltf_io.py`.

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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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'"
    ]
   },
   {
@@ -408,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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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'"
    ]
   },
   {
@@ -192,7 +195,7 @@
    "outputs": [],
    "source": [
     "# Load the dolphin mesh.\n",
-    "trg_obj = os.path.join('dolphin.obj')"
+    "trg_obj = 'dolphin.obj'"
    ]
   },
   {
@@ -247,7 +250,7 @@
     "id": "dYWDl4VGWHRK"
    },
    "source": [
-    "###  Visualize the source and target meshes"
+    "## 2. Visualize the source and target meshes"
    ]
   },
   {
@@ -265,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",
@@ -485,7 +488,7 @@
     "final_verts = final_verts * scale + center\n",
     "\n",
     "# Store the predicted mesh using save_obj\n",
-    "final_obj = os.path.join('./', 'final_model.obj')\n",
+    "final_obj = 'final_model.obj'\n",
     "save_obj(final_obj, final_verts, final_faces)"
    ]
   },

docs/tutorials/fit_simple_neural_radiance_field.ipynb

@@ -50,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-        "        !tar xzf 1.10.0.tar.gz\n",
-        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-        "        !tar xzf 1.10.0.tar.gz\n",
-        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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,28 +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.\") 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",
-    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
-    "        !tar xzf 1.10.0.tar.gz\n",
-    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\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/README.md

@@ -26,5 +26,6 @@ version_str="".join([
     torch.version.cuda.replace(".",""),
     f"_pyt{pyt_version_str}"
 ])
+!pip install fvcore iopath
 !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
 ```

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/README.md

@@ -212,9 +212,7 @@ from pytorch3d.implicitron.tools.config import registry
 class XRayRenderer(BaseRenderer, torch.nn.Module):
     n_pts_per_ray: int = 64
 
-    # if there are other base classes, make sure to call `super().__init__()` explicitly
     def __post_init__(self):
-        super().__init__()
         # custom initialization
 
     def forward(
@@ -250,7 +248,7 @@ The main object for this trainer loop is `Experiment`. It has four top-level rep
 * `data_source`: This is a `DataSourceBase` which defaults to `ImplicitronDataSource`.
 It constructs the data sets and dataloaders.
 * `model_factory`: This is a `ModelFactoryBase` which defaults to `ImplicitronModelFactory`.
-It constructs the model, which is usually an instance of implicitron's main `GenericModel` class, and can load its weights from a checkpoint.
+It constructs the model, which is usually an instance of `OverfitModel` (for NeRF-style training with overfitting to one scene) or `GenericModel` (that is able to generalize to multiple scenes by NeRFormer-style conditioning on other scene views), and can load its weights from a checkpoint.
 * `optimizer_factory`: This is an `OptimizerFactoryBase` which defaults to `ImplicitronOptimizerFactory`.
 It constructs the optimizer and can load its weights from a checkpoint.
 * `training_loop`: This is a `TrainingLoopBase` which defaults to `ImplicitronTrainingLoop` and defines the main training loop.
@@ -294,6 +292,43 @@ model_GenericModel_args: GenericModel
         ╘== ReductionFeatureAggregator
 ```
 
+Here is the class structure of OverfitModel:
+
+```
+model_OverfitModel_args: OverfitModel
+└-- raysampler_*_args: RaySampler
+    ╘== AdaptiveRaysampler
+    ╘== NearFarRaysampler
+└-- renderer_*_args: BaseRenderer
+    ╘== MultiPassEmissionAbsorptionRenderer
+    ╘== LSTMRenderer
+    ╘== SignedDistanceFunctionRenderer
+        └-- ray_tracer_args: RayTracing
+        └-- ray_normal_coloring_network_args: RayNormalColoringNetwork
+└-- implicit_function_*_args: ImplicitFunctionBase
+    ╘== NeuralRadianceFieldImplicitFunction
+    ╘== SRNImplicitFunction
+        └-- raymarch_function_args: SRNRaymarchFunction
+        └-- pixel_generator_args: SRNPixelGenerator
+    ╘== SRNHyperNetImplicitFunction
+        └-- hypernet_args: SRNRaymarchHyperNet
+        └-- pixel_generator_args: SRNPixelGenerator
+    ╘== IdrFeatureField
+└-- coarse_implicit_function_*_args: ImplicitFunctionBase
+    ╘== NeuralRadianceFieldImplicitFunction
+    ╘== SRNImplicitFunction
+        └-- raymarch_function_args: SRNRaymarchFunction
+        └-- pixel_generator_args: SRNPixelGenerator
+    ╘== SRNHyperNetImplicitFunction
+        └-- hypernet_args: SRNRaymarchHyperNet
+        └-- pixel_generator_args: SRNPixelGenerator
+    ╘== IdrFeatureField
+```
+
+OverfitModel has been introduced to create a simple class to disantagle Nerfs which the overfit pattern
+from the GenericModel.
+
+
 Please look at the annotations of the respective classes or functions for the lists of hyperparameters.
 `tests/experiment.yaml` shows every possible option if you have no user-defined classes.
 

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/configs/overfit_base.yaml

@@ -0,0 +1,79 @@
+defaults:
+- default_config
+- _self_
+exp_dir: ./data/exps/overfit_base/
+training_loop_ImplicitronTrainingLoop_args:
+  visdom_port: 8097
+  visualize_interval: 0
+  max_epochs: 1000
+data_source_ImplicitronDataSource_args:
+  data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
+  dataset_map_provider_class_type: JsonIndexDatasetMapProvider
+  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
+    dataset_length_train: 1000
+    dataset_length_val: 1
+    num_workers: 8
+  dataset_map_provider_JsonIndexDatasetMapProvider_args:
+    dataset_root: ${oc.env:CO3D_DATASET_ROOT}
+    n_frames_per_sequence: -1
+    test_on_train: true
+    test_restrict_sequence_id: 0
+    dataset_JsonIndexDataset_args:
+      load_point_clouds: false
+      mask_depths: false
+      mask_images: false
+model_factory_ImplicitronModelFactory_args:
+  model_class_type: "OverfitModel"
+  model_OverfitModel_args:
+    loss_weights:
+      loss_mask_bce: 1.0
+      loss_prev_stage_mask_bce: 1.0
+      loss_autodecoder_norm: 0.01
+      loss_rgb_mse: 1.0
+      loss_prev_stage_rgb_mse: 1.0
+    output_rasterized_mc: false
+    chunk_size_grid: 102400
+    render_image_height: 400
+    render_image_width: 400
+    share_implicit_function_across_passes: false
+    implicit_function_class_type: "NeuralRadianceFieldImplicitFunction"
+    implicit_function_NeuralRadianceFieldImplicitFunction_args:
+      n_harmonic_functions_xyz: 10
+      n_harmonic_functions_dir: 4
+      n_hidden_neurons_xyz: 256
+      n_hidden_neurons_dir: 128
+      n_layers_xyz: 8
+      append_xyz:
+      - 5
+    coarse_implicit_function_class_type: "NeuralRadianceFieldImplicitFunction"
+    coarse_implicit_function_NeuralRadianceFieldImplicitFunction_args:
+      n_harmonic_functions_xyz: 10
+      n_harmonic_functions_dir: 4
+      n_hidden_neurons_xyz: 256
+      n_hidden_neurons_dir: 128
+      n_layers_xyz: 8
+      append_xyz:
+      - 5
+    raysampler_AdaptiveRaySampler_args:
+      n_rays_per_image_sampled_from_mask: 1024
+      scene_extent: 8.0
+      n_pts_per_ray_training: 64
+      n_pts_per_ray_evaluation: 64
+      stratified_point_sampling_training: true
+      stratified_point_sampling_evaluation: false
+    renderer_MultiPassEmissionAbsorptionRenderer_args:
+      n_pts_per_ray_fine_training: 64
+      n_pts_per_ray_fine_evaluation: 64
+      append_coarse_samples_to_fine: true
+      density_noise_std_train: 1.0
+optimizer_factory_ImplicitronOptimizerFactory_args:
+  breed: Adam
+  weight_decay: 0.0
+  lr_policy: MultiStepLR
+  multistep_lr_milestones: []
+  lr: 0.0005
+  gamma: 0.1
+  momentum: 0.9
+  betas:
+  - 0.9
+  - 0.999

projects/implicitron_trainer/configs/overfit_singleseq_base.yaml

@@ -0,0 +1,42 @@
+defaults:
+- overfit_base
+- _self_
+data_source_ImplicitronDataSource_args:
+  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
+    batch_size: 1
+    dataset_length_train: 1000
+    dataset_length_val: 1
+    num_workers: 8
+  dataset_map_provider_JsonIndexDatasetMapProvider_args:
+    assert_single_seq: true
+    n_frames_per_sequence: -1
+    test_restrict_sequence_id: 0
+    test_on_train: false
+model_factory_ImplicitronModelFactory_args:
+  model_class_type: "OverfitModel"
+  model_OverfitModel_args:
+    render_image_height: 800
+    render_image_width: 800
+    log_vars:
+    - loss_rgb_psnr_fg
+    - loss_rgb_psnr
+    - loss_eikonal
+    - loss_prev_stage_rgb_psnr
+    - loss_mask_bce
+    - loss_prev_stage_mask_bce
+    - loss_rgb_mse
+    - loss_prev_stage_rgb_mse
+    - loss_depth_abs
+    - loss_depth_abs_fg
+    - loss_kl
+    - loss_mask_neg_iou
+    - objective
+    - epoch
+    - sec/it
+optimizer_factory_ImplicitronOptimizerFactory_args:
+  lr: 0.0005
+  multistep_lr_milestones:
+  - 200
+  - 300
+training_loop_ImplicitronTrainingLoop_args:
+  max_epochs: 400

projects/implicitron_trainer/configs/overfit_singleseq_nerf_blender.yaml

@@ -0,0 +1,56 @@
+defaults:
+- overfit_singleseq_base
+- _self_
+exp_dir: "./data/overfit_nerf_blender_repro/${oc.env:BLENDER_SINGLESEQ_CLASS}"
+data_source_ImplicitronDataSource_args:
+  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
+    dataset_length_train: 100
+  dataset_map_provider_class_type: BlenderDatasetMapProvider
+  dataset_map_provider_BlenderDatasetMapProvider_args:
+    base_dir: ${oc.env:BLENDER_DATASET_ROOT}/${oc.env:BLENDER_SINGLESEQ_CLASS}
+    n_known_frames_for_test: null
+    object_name: ${oc.env:BLENDER_SINGLESEQ_CLASS}
+    path_manager_factory_class_type: PathManagerFactory
+    path_manager_factory_PathManagerFactory_args:
+      silence_logs: true
+
+model_factory_ImplicitronModelFactory_args:
+  model_class_type: "OverfitModel"
+  model_OverfitModel_args:
+    mask_images: false
+    raysampler_class_type: AdaptiveRaySampler
+    raysampler_AdaptiveRaySampler_args:
+      n_pts_per_ray_training: 64
+      n_pts_per_ray_evaluation: 64
+      n_rays_per_image_sampled_from_mask: 4096
+      stratified_point_sampling_training: true
+      stratified_point_sampling_evaluation: false
+      scene_extent: 2.0
+      scene_center:
+      - 0.0
+      - 0.0
+      - 0.0
+    renderer_MultiPassEmissionAbsorptionRenderer_args:
+      density_noise_std_train: 0.0
+      n_pts_per_ray_fine_training: 128
+      n_pts_per_ray_fine_evaluation: 128
+      raymarcher_EmissionAbsorptionRaymarcher_args:
+        blend_output: false
+    loss_weights:
+      loss_rgb_mse: 1.0
+      loss_prev_stage_rgb_mse: 1.0
+      loss_mask_bce: 0.0
+      loss_prev_stage_mask_bce: 0.0
+      loss_autodecoder_norm: 0.00
+
+optimizer_factory_ImplicitronOptimizerFactory_args:
+  exponential_lr_step_size: 3001
+  lr_policy: LinearExponential
+  linear_exponential_lr_milestone: 200
+
+training_loop_ImplicitronTrainingLoop_args:
+  max_epochs: 6000
+  metric_print_interval: 10
+  store_checkpoints_purge: 3
+  test_when_finished: true
+  validation_interval: 100

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.
 
@@ -59,7 +61,7 @@ from pytorch3d.implicitron.dataset.data_source import (
     DataSourceBase,
     ImplicitronDataSource,
 )
-from pytorch3d.implicitron.models.generic_model import ImplicitronModelBase
+from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
 
 from pytorch3d.implicitron.models.renderer.multipass_ea import (
     MultiPassEmissionAbsorptionRenderer,
@@ -207,12 +209,6 @@ class Experiment(Configurable):  # pyre-ignore: 13
                 val_loader,
             ) = accelerator.prepare(model, optimizer, train_loader, val_loader)
 
-        # pyre-fixme[16]: Optional type has no attribute `is_multisequence`.
-        if not self.training_loop.evaluator.is_multisequence:
-            all_train_cameras = self.data_source.all_train_cameras
-        else:
-            all_train_cameras = None
-
         # Enter the main training loop.
         self.training_loop.run(
             train_loader=train_loader,
@@ -223,7 +219,6 @@ class Experiment(Configurable):  # pyre-ignore: 13
             model=model,
             optimizer=optimizer,
             scheduler=scheduler,
-            all_train_cameras=all_train_cameras,
             accelerator=accelerator,
             device=device,
             exp_dir=self.exp_dir,

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,
         *,
@@ -122,7 +125,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
         optimizer: torch.optim.Optimizer,
         scheduler: Any,
         accelerator: Optional[Accelerator],
-        all_train_cameras: Optional[CamerasBase],
         device: torch.device,
         exp_dir: str,
         stats: Stats,
@@ -142,7 +144,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
             if test_loader is not None:
                 # pyre-fixme[16]: `Optional` has no attribute `run`.
                 self.evaluator.run(
-                    all_train_cameras=all_train_cameras,
                     dataloader=test_loader,
                     device=device,
                     dump_to_json=True,
@@ -200,7 +201,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
                     and epoch % self.test_interval == 0
                 ):
                     self.evaluator.run(
-                        all_train_cameras=all_train_cameras,
                         device=device,
                         dataloader=test_loader,
                         model=model,
@@ -217,7 +217,6 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
         if self.test_when_finished:
             if test_loader is not None:
                 self.evaluator.run(
-                    all_train_cameras=all_train_cameras,
                     device=device,
                     dump_to_json=True,
                     epoch=stats.epoch,
@@ -260,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,
         )
@@ -386,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 (
@@ -396,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

@@ -103,8 +103,10 @@ data_source_ImplicitronDataSource_args:
     num_views: 40
     data_file: null
     azimuth_range: 180.0
+    distance: 2.7
     resolution: 128
     use_point_light: true
+    gpu_idx: 0
     path_manager_factory_class_type: PathManagerFactory
     path_manager_factory_PathManagerFactory_args:
       silence_logs: true
@@ -127,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
@@ -201,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
@@ -213,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:
@@ -232,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:
@@ -344,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
@@ -355,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
@@ -559,6 +580,631 @@ model_factory_ImplicitronModelFactory_args:
           use_xavier_init: true
     view_metrics_ViewMetrics_args: {}
     regularization_metrics_RegularizationMetrics_args: {}
+  model_OverfitModel_args:
+    log_vars:
+    - loss_rgb_psnr_fg
+    - loss_rgb_psnr
+    - loss_rgb_mse
+    - loss_rgb_huber
+    - loss_depth_abs
+    - loss_depth_abs_fg
+    - loss_mask_neg_iou
+    - loss_mask_bce
+    - loss_mask_beta_prior
+    - loss_eikonal
+    - loss_density_tv
+    - loss_depth_neg_penalty
+    - loss_autodecoder_norm
+    - loss_prev_stage_rgb_mse
+    - loss_prev_stage_rgb_psnr_fg
+    - loss_prev_stage_rgb_psnr
+    - loss_prev_stage_mask_bce
+    - objective
+    - epoch
+    - sec/it
+    mask_images: true
+    mask_depths: true
+    render_image_width: 400
+    render_image_height: 400
+    mask_threshold: 0.5
+    output_rasterized_mc: false
+    bg_color:
+    - 0.0
+    - 0.0
+    - 0.0
+    chunk_size_grid: 4096
+    render_features_dimensions: 3
+    tqdm_trigger_threshold: 16
+    n_train_target_views: 1
+    sampling_mode_training: mask_sample
+    sampling_mode_evaluation: full_grid
+    global_encoder_class_type: null
+    raysampler_class_type: AdaptiveRaySampler
+    renderer_class_type: MultiPassEmissionAbsorptionRenderer
+    share_implicit_function_across_passes: false
+    implicit_function_class_type: NeuralRadianceFieldImplicitFunction
+    coarse_implicit_function_class_type: null
+    view_metrics_class_type: ViewMetrics
+    regularization_metrics_class_type: RegularizationMetrics
+    loss_weights:
+      loss_rgb_mse: 1.0
+      loss_prev_stage_rgb_mse: 1.0
+      loss_mask_bce: 0.0
+      loss_prev_stage_mask_bce: 0.0
+    global_encoder_HarmonicTimeEncoder_args:
+      n_harmonic_functions: 10
+      append_input: true
+      time_divisor: 1.0
+    global_encoder_SequenceAutodecoder_args:
+      autodecoder_args:
+        encoding_dim: 0
+        n_instances: 1
+        init_scale: 1.0
+        ignore_input: false
+    raysampler_AdaptiveRaySampler_args:
+      n_pts_per_ray_training: 64
+      n_pts_per_ray_evaluation: 64
+      n_rays_per_image_sampled_from_mask: 1024
+      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
+      - 0.0
+      - 0.0
+    raysampler_NearFarRaySampler_args:
+      n_pts_per_ray_training: 64
+      n_pts_per_ray_evaluation: 64
+      n_rays_per_image_sampled_from_mask: 1024
+      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:
+      num_raymarch_steps: 10
+      init_depth: 17.0
+      init_depth_noise_std: 0.0005
+      hidden_size: 16
+      n_feature_channels: 256
+      bg_color: null
+      verbose: false
+    renderer_MultiPassEmissionAbsorptionRenderer_args:
+      raymarcher_class_type: EmissionAbsorptionRaymarcher
+      n_pts_per_ray_fine_training: 64
+      n_pts_per_ray_fine_evaluation: 64
+      stratified_sampling_coarse_training: true
+      stratified_sampling_coarse_evaluation: false
+      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:
+        - 0.0
+        replicate_last_interval: false
+        background_opacity: 0.0
+        density_relu: true
+        blend_output: false
+      raymarcher_EmissionAbsorptionRaymarcher_args:
+        surface_thickness: 1
+        bg_color:
+        - 0.0
+        replicate_last_interval: false
+        background_opacity: 10000000000.0
+        density_relu: true
+        blend_output: false
+    renderer_SignedDistanceFunctionRenderer_args:
+      ray_normal_coloring_network_args:
+        feature_vector_size: 3
+        mode: idr
+        d_in: 9
+        d_out: 3
+        dims:
+        - 512
+        - 512
+        - 512
+        - 512
+        weight_norm: true
+        n_harmonic_functions_dir: 0
+        pooled_feature_dim: 0
+      bg_color:
+      - 0.0
+      soft_mask_alpha: 50.0
+      ray_tracer_args:
+        sdf_threshold: 5.0e-05
+        line_search_step: 0.5
+        line_step_iters: 1
+        sphere_tracing_iters: 10
+        n_steps: 100
+        n_secant_steps: 8
+    implicit_function_IdrFeatureField_args:
+      d_in: 3
+      d_out: 1
+      dims:
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      geometric_init: true
+      bias: 1.0
+      skip_in: []
+      weight_norm: true
+      n_harmonic_functions_xyz: 0
+      pooled_feature_dim: 0
+    implicit_function_NeRFormerImplicitFunction_args:
+      n_harmonic_functions_xyz: 10
+      n_harmonic_functions_dir: 4
+      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
+      append_xyz:
+      - 1
+    implicit_function_NeuralRadianceFieldImplicitFunction_args:
+      n_harmonic_functions_xyz: 10
+      n_harmonic_functions_dir: 4
+      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
+      append_xyz:
+      - 5
+    implicit_function_SRNHyperNetImplicitFunction_args:
+      latent_dim_hypernet: 0
+      hypernet_args:
+        n_harmonic_functions: 3
+        n_hidden_units: 256
+        n_layers: 2
+        n_hidden_units_hypernet: 256
+        n_layers_hypernet: 1
+        in_features: 3
+        out_features: 256
+        xyz_in_camera_coords: false
+      pixel_generator_args:
+        n_harmonic_functions: 4
+        n_hidden_units: 256
+        n_hidden_units_color: 128
+        n_layers: 2
+        in_features: 256
+        out_features: 3
+        ray_dir_in_camera_coords: false
+    implicit_function_SRNImplicitFunction_args:
+      raymarch_function_args:
+        n_harmonic_functions: 3
+        n_hidden_units: 256
+        n_layers: 2
+        in_features: 3
+        out_features: 256
+        xyz_in_camera_coords: false
+        raymarch_function: null
+      pixel_generator_args:
+        n_harmonic_functions: 4
+        n_hidden_units: 256
+        n_hidden_units_color: 128
+        n_layers: 2
+        in_features: 256
+        out_features: 3
+        ray_dir_in_camera_coords: false
+    implicit_function_VoxelGridImplicitFunction_args:
+      harmonic_embedder_xyz_density_args:
+        n_harmonic_functions: 6
+        omega_0: 1.0
+        logspace: true
+        append_input: true
+      harmonic_embedder_xyz_color_args:
+        n_harmonic_functions: 6
+        omega_0: 1.0
+        logspace: true
+        append_input: true
+      harmonic_embedder_dir_color_args:
+        n_harmonic_functions: 6
+        omega_0: 1.0
+        logspace: true
+        append_input: true
+      decoder_density_class_type: MLPDecoder
+      decoder_color_class_type: MLPDecoder
+      use_multiple_streams: true
+      xyz_ray_dir_in_camera_coords: false
+      scaffold_calculating_epochs: []
+      scaffold_resolution:
+      - 128
+      - 128
+      - 128
+      scaffold_empty_space_threshold: 0.001
+      scaffold_occupancy_chunk_size: -1
+      scaffold_max_pool_kernel_size: 3
+      scaffold_filter_points: true
+      volume_cropping_epochs: []
+      voxel_grid_density_args:
+        voxel_grid_class_type: FullResolutionVoxelGrid
+        extents:
+        - 2.0
+        - 2.0
+        - 2.0
+        translation:
+        - 0.0
+        - 0.0
+        - 0.0
+        init_std: 0.1
+        init_mean: 0.0
+        hold_voxel_grid_as_parameters: true
+        param_groups: {}
+        voxel_grid_CPFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: 24
+          basis_matrix: true
+        voxel_grid_FullResolutionVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+        voxel_grid_VMFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: null
+          distribution_of_components: null
+          basis_matrix: true
+      voxel_grid_color_args:
+        voxel_grid_class_type: FullResolutionVoxelGrid
+        extents:
+        - 2.0
+        - 2.0
+        - 2.0
+        translation:
+        - 0.0
+        - 0.0
+        - 0.0
+        init_std: 0.1
+        init_mean: 0.0
+        hold_voxel_grid_as_parameters: true
+        param_groups: {}
+        voxel_grid_CPFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: 24
+          basis_matrix: true
+        voxel_grid_FullResolutionVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+        voxel_grid_VMFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: null
+          distribution_of_components: null
+          basis_matrix: true
+      decoder_density_ElementwiseDecoder_args:
+        scale: 1.0
+        shift: 0.0
+        operation: IDENTITY
+      decoder_density_MLPDecoder_args:
+        param_groups: {}
+        network_args:
+          n_layers: 8
+          output_dim: 256
+          skip_dim: 39
+          hidden_dim: 256
+          input_skips:
+          - 5
+          skip_affine_trans: false
+          last_layer_bias_init: null
+          last_activation: RELU
+          use_xavier_init: true
+      decoder_color_ElementwiseDecoder_args:
+        scale: 1.0
+        shift: 0.0
+        operation: IDENTITY
+      decoder_color_MLPDecoder_args:
+        param_groups: {}
+        network_args:
+          n_layers: 8
+          output_dim: 256
+          skip_dim: 39
+          hidden_dim: 256
+          input_skips:
+          - 5
+          skip_affine_trans: false
+          last_layer_bias_init: null
+          last_activation: RELU
+          use_xavier_init: true
+    coarse_implicit_function_IdrFeatureField_args:
+      d_in: 3
+      d_out: 1
+      dims:
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      - 512
+      geometric_init: true
+      bias: 1.0
+      skip_in: []
+      weight_norm: true
+      n_harmonic_functions_xyz: 0
+      pooled_feature_dim: 0
+    coarse_implicit_function_NeRFormerImplicitFunction_args:
+      n_harmonic_functions_xyz: 10
+      n_harmonic_functions_dir: 4
+      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
+      append_xyz:
+      - 1
+    coarse_implicit_function_NeuralRadianceFieldImplicitFunction_args:
+      n_harmonic_functions_xyz: 10
+      n_harmonic_functions_dir: 4
+      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
+      append_xyz:
+      - 5
+    coarse_implicit_function_SRNHyperNetImplicitFunction_args:
+      latent_dim_hypernet: 0
+      hypernet_args:
+        n_harmonic_functions: 3
+        n_hidden_units: 256
+        n_layers: 2
+        n_hidden_units_hypernet: 256
+        n_layers_hypernet: 1
+        in_features: 3
+        out_features: 256
+        xyz_in_camera_coords: false
+      pixel_generator_args:
+        n_harmonic_functions: 4
+        n_hidden_units: 256
+        n_hidden_units_color: 128
+        n_layers: 2
+        in_features: 256
+        out_features: 3
+        ray_dir_in_camera_coords: false
+    coarse_implicit_function_SRNImplicitFunction_args:
+      raymarch_function_args:
+        n_harmonic_functions: 3
+        n_hidden_units: 256
+        n_layers: 2
+        in_features: 3
+        out_features: 256
+        xyz_in_camera_coords: false
+        raymarch_function: null
+      pixel_generator_args:
+        n_harmonic_functions: 4
+        n_hidden_units: 256
+        n_hidden_units_color: 128
+        n_layers: 2
+        in_features: 256
+        out_features: 3
+        ray_dir_in_camera_coords: false
+    coarse_implicit_function_VoxelGridImplicitFunction_args:
+      harmonic_embedder_xyz_density_args:
+        n_harmonic_functions: 6
+        omega_0: 1.0
+        logspace: true
+        append_input: true
+      harmonic_embedder_xyz_color_args:
+        n_harmonic_functions: 6
+        omega_0: 1.0
+        logspace: true
+        append_input: true
+      harmonic_embedder_dir_color_args:
+        n_harmonic_functions: 6
+        omega_0: 1.0
+        logspace: true
+        append_input: true
+      decoder_density_class_type: MLPDecoder
+      decoder_color_class_type: MLPDecoder
+      use_multiple_streams: true
+      xyz_ray_dir_in_camera_coords: false
+      scaffold_calculating_epochs: []
+      scaffold_resolution:
+      - 128
+      - 128
+      - 128
+      scaffold_empty_space_threshold: 0.001
+      scaffold_occupancy_chunk_size: -1
+      scaffold_max_pool_kernel_size: 3
+      scaffold_filter_points: true
+      volume_cropping_epochs: []
+      voxel_grid_density_args:
+        voxel_grid_class_type: FullResolutionVoxelGrid
+        extents:
+        - 2.0
+        - 2.0
+        - 2.0
+        translation:
+        - 0.0
+        - 0.0
+        - 0.0
+        init_std: 0.1
+        init_mean: 0.0
+        hold_voxel_grid_as_parameters: true
+        param_groups: {}
+        voxel_grid_CPFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: 24
+          basis_matrix: true
+        voxel_grid_FullResolutionVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+        voxel_grid_VMFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: null
+          distribution_of_components: null
+          basis_matrix: true
+      voxel_grid_color_args:
+        voxel_grid_class_type: FullResolutionVoxelGrid
+        extents:
+        - 2.0
+        - 2.0
+        - 2.0
+        translation:
+        - 0.0
+        - 0.0
+        - 0.0
+        init_std: 0.1
+        init_mean: 0.0
+        hold_voxel_grid_as_parameters: true
+        param_groups: {}
+        voxel_grid_CPFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: 24
+          basis_matrix: true
+        voxel_grid_FullResolutionVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+        voxel_grid_VMFactorizedVoxelGrid_args:
+          align_corners: true
+          padding: zeros
+          mode: bilinear
+          n_features: 1
+          resolution_changes:
+            0:
+            - 128
+            - 128
+            - 128
+          n_components: null
+          distribution_of_components: null
+          basis_matrix: true
+      decoder_density_ElementwiseDecoder_args:
+        scale: 1.0
+        shift: 0.0
+        operation: IDENTITY
+      decoder_density_MLPDecoder_args:
+        param_groups: {}
+        network_args:
+          n_layers: 8
+          output_dim: 256
+          skip_dim: 39
+          hidden_dim: 256
+          input_skips:
+          - 5
+          skip_affine_trans: false
+          last_layer_bias_init: null
+          last_activation: RELU
+          use_xavier_init: true
+      decoder_color_ElementwiseDecoder_args:
+        scale: 1.0
+        shift: 0.0
+        operation: IDENTITY
+      decoder_color_MLPDecoder_args:
+        param_groups: {}
+        network_args:
+          n_layers: 8
+          output_dim: 256
+          skip_dim: 39
+          hidden_dim: 256
+          input_skips:
+          - 5
+          skip_affine_trans: false
+          last_layer_bias_init: null
+          last_activation: RELU
+          use_xavier_init: true
+    view_metrics_ViewMetrics_args: {}
+    regularization_metrics_RegularizationMetrics_args: {}
 optimizer_factory_ImplicitronOptimizerFactory_args:
   betas:
   - 0.9

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)
@@ -141,7 +150,11 @@ class TestExperiment(unittest.TestCase):
         # Check that all the pre-prepared configs are valid.
         config_files = []
 
-        for pattern in ("repro_singleseq*.yaml", "repro_multiseq*.yaml"):
+        for pattern in (
+            "repro_singleseq*.yaml",
+            "repro_multiseq*.yaml",
+            "overfit_singleseq*.yaml",
+        ):
             config_files.extend(
                 [
                     f

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:
 
@@ -39,6 +41,7 @@ def visualize_reconstruction(
     visdom_server: str = "http://127.0.0.1",
     visdom_port: int = 8097,
     visdom_env: Optional[str] = None,
+    **render_flyaround_kwargs,
 ) -> None:
     """
     Given an `exp_dir` containing a trained Implicitron model, generates videos consisting
@@ -60,6 +63,8 @@ def visualize_reconstruction(
         visdom_server: The address of the visdom server.
         visdom_port: The port of the visdom server.
         visdom_env: If set, defines a custom name for the visdom environment.
+        render_flyaround_kwargs: Keyword arguments passed to the invoked `render_flyaround`
+            function (see `pytorch3d.implicitron.models.visualization.render_flyaround`).
     """
 
     # In case an output directory is specified use it. If no output_directory
@@ -115,20 +120,22 @@ def visualize_reconstruction(
     # iterate over the sequences in the dataset
     for sequence_name in dataset.sequence_names():
         with torch.no_grad():
-            render_flyaround(
-                dataset=dataset,
-                sequence_name=sequence_name,
-                model=model,
-                output_video_path=os.path.join(output_directory, "video"),
-                n_source_views=n_source_views,
-                visdom_show_preds=visdom_show_preds,
-                n_flyaround_poses=n_eval_cameras,
-                visdom_server=visdom_server,
-                visdom_port=visdom_port,
-                visdom_environment=visdom_env,
-                video_resize=video_size,
-                device=device,
-            )
+            render_kwargs = {
+                "dataset": dataset,
+                "sequence_name": sequence_name,
+                "model": model,
+                "output_video_path": os.path.join(output_directory, "video"),
+                "n_source_views": n_source_views,
+                "visdom_show_preds": visdom_show_preds,
+                "n_flyaround_poses": n_eval_cameras,
+                "visdom_server": visdom_server,
+                "visdom_port": visdom_port,
+                "visdom_environment": visdom_env,
+                "video_resize": video_size,
+                "device": device,
+                **render_flyaround_kwargs,
+            }
+            render_flyaround(**render_kwargs)
 
 
 enable_get_default_args(visualize_reconstruction)

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.2"
+# 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
@@ -57,19 +58,3 @@ def get_device(x, device: Optional[Device] = None) -> torch.device:
 
     # Default device is cpu
     return torch.device("cpu")
-
-
-# Provide get_origin and get_args even in Python 3.7.
-
-if sys.version_info >= (3, 8, 0):
-    from typing import get_args, get_origin
-elif sys.version_info >= (3, 7, 0):
-
-    def get_origin(cls):  # pragma: no cover
-        return getattr(cls, "__origin__", None)
-
-    def get_args(cls):  # pragma: no cover
-        return getattr(cls, "__args__", None)
-
-else:
-    raise ImportError("This module requires Python 3.7+")

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/points_to_volumes/points_to_volumes.h

@@ -7,6 +7,7 @@
  */
 
 #pragma once
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <cstdio>
 #include <tuple>
@@ -96,6 +97,8 @@ inline void PointsToVolumesForward(
         point_weight,
         align_corners,
         splat);
+    torch::autograd::increment_version(volume_features);
+    torch::autograd::increment_version(volume_densities);
     return;
 #else
     AT_ERROR("Not compiled with GPU support.");

pytorch3d/csrc/points_to_volumes/points_to_volumes_cpu.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <algorithm>
 #include <cmath>
@@ -148,6 +149,8 @@ void PointsToVolumesForwardCpu(
       }
     }
   }
+  torch::autograd::increment_version(volume_features);
+  torch::autograd::increment_version(volume_densities);
 }
 
 // With nearest, the only smooth dependence is that volume features

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/csrc/sample_pdf/sample_pdf.h

@@ -7,6 +7,7 @@
  */
 
 #pragma once
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <cstdio>
 #include <tuple>
@@ -63,6 +64,7 @@ inline void SamplePdf(
 #ifdef WITH_CUDA
     CHECK_CUDA(weights);
     CHECK_CONTIGUOUS_CUDA(outputs);
+    torch::autograd::increment_version(outputs);
     SamplePdfCuda(bins, weights, outputs, eps);
     return;
 #else

pytorch3d/csrc/sample_pdf/sample_pdf_cpu.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <torch/csrc/autograd/VariableTypeUtils.h>
 #include <torch/extension.h>
 #include <algorithm>
 #include <thread>
@@ -137,4 +138,5 @@ void SamplePdfCpu(
   for (auto&& thread : threads) {
     thread.join();
   }
+  torch::autograd::increment_version(outputs);
 }

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
@@ -12,14 +14,15 @@ import torch
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
 from torch.utils.data import (
     BatchSampler,
-    ChainDataset,
+    ConcatDataset,
     DataLoader,
     RandomSampler,
     Sampler,
 )
 
-from .dataset_base import DatasetBase, FrameData
+from .dataset_base import DatasetBase
 from .dataset_map_provider import DatasetMap
+from .frame_data import FrameData
 from .scene_batch_sampler import SceneBatchSampler
 from .utils import is_known_frame_scalar
 
@@ -482,7 +485,7 @@ class SequenceDataLoaderMapProvider(DataLoaderMapProviderBase):
             num_batches=num_batches,
         )
         return DataLoader(
-            ChainDataset([dataset, train_dataset]),
+            ConcatDataset([dataset, train_dataset]),
             batch_sampler=sampler,
             **data_loader_kwargs,
         )

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 (
@@ -13,13 +15,8 @@ from pytorch3d.implicitron.tools.config import (
 )
 from pytorch3d.renderer.cameras import CamerasBase
 
-from .blender_dataset_map_provider import BlenderDatasetMapProvider  # noqa
 from .data_loader_map_provider import DataLoaderMap, DataLoaderMapProviderBase
 from .dataset_map_provider import DatasetMap, DatasetMapProviderBase
-from .json_index_dataset_map_provider import JsonIndexDatasetMapProvider  # noqa
-from .json_index_dataset_map_provider_v2 import JsonIndexDatasetMapProviderV2  # noqa
-from .llff_dataset_map_provider import LlffDatasetMapProvider  # noqa
-from .rendered_mesh_dataset_map_provider import RenderedMeshDatasetMapProvider  # noqa
 
 
 class DataSourceBase(ReplaceableBase):
@@ -34,6 +31,7 @@ class DataSourceBase(ReplaceableBase):
     @property
     def all_train_cameras(self) -> Optional[CamerasBase]:
         """
+        DEPRECATED! The property will be removed in future versions.
         If the data is all for a single scene, a list
         of the known training cameras for that scene, which is
         used for evaluating the viewpoint difficulty of the
@@ -59,6 +57,36 @@ class ImplicitronDataSource(DataSourceBase):  # pyre-ignore[13]
     data_loader_map_provider: DataLoaderMapProviderBase
     data_loader_map_provider_class_type: str = "SequenceDataLoaderMapProvider"
 
+    @classmethod
+    def pre_expand(cls) -> None:
+        # use try/finally to bypass cinder's lazy imports
+        try:
+            from .blender_dataset_map_provider import (  # noqa: F401
+                BlenderDatasetMapProvider,
+            )
+            from .json_index_dataset_map_provider import (  # noqa: F401
+                JsonIndexDatasetMapProvider,
+            )
+            from .json_index_dataset_map_provider_v2 import (  # noqa: F401
+                JsonIndexDatasetMapProviderV2,
+            )
+            from .llff_dataset_map_provider import LlffDatasetMapProvider  # noqa: F401
+            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
+
     def __post_init__(self):
         run_auto_creation(self)
         self._all_train_cameras_cache: Optional[Tuple[Optional[CamerasBase]]] = None
@@ -70,6 +98,9 @@ class ImplicitronDataSource(DataSourceBase):  # pyre-ignore[13]
 
     @property
     def all_train_cameras(self) -> Optional[CamerasBase]:
+        """
+        DEPRECATED! The property will be removed in future versions.
+        """
         if self._all_train_cameras_cache is None:  # pyre-ignore[16]
             all_train_cameras = self.dataset_map_provider.get_all_train_cameras()
             self._all_train_cameras_cache = (all_train_cameras,)

pytorch3d/implicitron/dataset/dataset_base.py

@@ -4,218 +4,30 @@
 # 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, field, fields
+from dataclasses import dataclass
 from typing import (
-    Any,
     ClassVar,
     Dict,
     Iterable,
     Iterator,
     List,
-    Mapping,
     Optional,
     Sequence,
     Tuple,
     Type,
-    Union,
 )
 
-import numpy as np
 import torch
-from pytorch3d.renderer.camera_utils import join_cameras_as_batch
-from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
-from pytorch3d.structures.pointclouds import join_pointclouds_as_batch, Pointclouds
 
-
-@dataclass
-class FrameData(Mapping[str, Any]):
-    """
-    A type of the elements returned by indexing the dataset object.
-    It can represent both individual frames and batches of thereof;
-    in this documentation, the sizes of tensors refer to single frames;
-    add the first batch dimension for the collation result.
-
-    Args:
-        frame_number: The number of the frame within its sequence.
-            0-based continuous integers.
-        sequence_name: The unique name of the frame's sequence.
-        sequence_category: The object category of the sequence.
-        frame_timestamp: The time elapsed since the start of a sequence in sec.
-        image_size_hw: The size of the image in pixels; (height, width) tensor
-                        of shape (2,).
-        image_path: The qualified path to the loaded image (with dataset_root).
-        image_rgb: A Tensor of shape `(3, H, W)` holding the RGB image
-            of the frame; elements are floats in [0, 1].
-        mask_crop: A binary mask of shape `(1, H, W)` denoting the valid image
-            regions. Regions can be invalid (mask_crop[i,j]=0) in case they
-            are a result of zero-padding of the image after cropping around
-            the object bounding box; elements are floats in {0.0, 1.0}.
-        depth_path: The qualified path to the frame's depth map.
-        depth_map: A float Tensor of shape `(1, H, W)` holding the depth map
-            of the frame; values correspond to distances from the camera;
-            use `depth_mask` and `mask_crop` to filter for valid pixels.
-        depth_mask: A binary mask of shape `(1, H, W)` denoting pixels of the
-            depth map that are valid for evaluation, they have been checked for
-            consistency across views; elements are floats in {0.0, 1.0}.
-        mask_path: A qualified path to the foreground probability mask.
-        fg_probability: A Tensor of `(1, H, W)` denoting the probability of the
-            pixels belonging to the captured object; elements are floats
-            in [0, 1].
-        bbox_xywh: The bounding box tightly enclosing the foreground object in the
-            format (x0, y0, width, height). The convention assumes that
-            `x0+width` and `y0+height` includes the boundary of the box.
-            I.e., to slice out the corresponding crop from an image tensor `I`
-            we execute `crop = I[..., y0:y0+height, x0:x0+width]`
-        crop_bbox_xywh: The bounding box denoting the boundaries of `image_rgb`
-            in the original image coordinates in the format (x0, y0, width, height).
-            The convention is the same as for `bbox_xywh`. `crop_bbox_xywh` differs
-            from `bbox_xywh` due to padding (which can happen e.g. due to
-            setting `JsonIndexDataset.box_crop_context > 0`)
-        camera: A PyTorch3D camera object corresponding the frame's viewpoint,
-            corrected for cropping if it happened.
-        camera_quality_score: The score proportional to the confidence of the
-            frame's camera estimation (the higher the more accurate).
-        point_cloud_quality_score: The score proportional to the accuracy of the
-            frame's sequence point cloud (the higher the more accurate).
-        sequence_point_cloud_path: The path to the sequence's point cloud.
-        sequence_point_cloud: A PyTorch3D Pointclouds object holding the
-            point cloud corresponding to the frame's sequence. When the object
-            represents a batch of frames, point clouds may be deduplicated;
-            see `sequence_point_cloud_idx`.
-        sequence_point_cloud_idx: Integer indices mapping frame indices to the
-            corresponding point clouds in `sequence_point_cloud`; to get the
-            corresponding point cloud to `image_rgb[i]`, use
-            `sequence_point_cloud[sequence_point_cloud_idx[i]]`.
-        frame_type: The type of the loaded frame specified in
-            `subset_lists_file`, if provided.
-        meta: A dict for storing additional frame information.
-    """
-
-    frame_number: Optional[torch.LongTensor]
-    sequence_name: Union[str, List[str]]
-    sequence_category: Union[str, List[str]]
-    frame_timestamp: Optional[torch.Tensor] = None
-    image_size_hw: Optional[torch.Tensor] = None
-    image_path: Union[str, List[str], None] = None
-    image_rgb: Optional[torch.Tensor] = None
-    # masks out padding added due to cropping the square bit
-    mask_crop: Optional[torch.Tensor] = None
-    depth_path: Union[str, List[str], None] = None
-    depth_map: Optional[torch.Tensor] = None
-    depth_mask: Optional[torch.Tensor] = None
-    mask_path: Union[str, List[str], None] = None
-    fg_probability: Optional[torch.Tensor] = None
-    bbox_xywh: Optional[torch.Tensor] = None
-    crop_bbox_xywh: Optional[torch.Tensor] = None
-    camera: Optional[PerspectiveCameras] = None
-    camera_quality_score: Optional[torch.Tensor] = None
-    point_cloud_quality_score: Optional[torch.Tensor] = None
-    sequence_point_cloud_path: Union[str, List[str], None] = None
-    sequence_point_cloud: Optional[Pointclouds] = None
-    sequence_point_cloud_idx: Optional[torch.Tensor] = None
-    frame_type: Union[str, List[str], None] = None  # known | unseen
-    meta: dict = field(default_factory=lambda: {})
-
-    def to(self, *args, **kwargs):
-        new_params = {}
-        for f in fields(self):
-            value = getattr(self, f.name)
-            if isinstance(value, (torch.Tensor, Pointclouds, CamerasBase)):
-                new_params[f.name] = value.to(*args, **kwargs)
-            else:
-                new_params[f.name] = value
-        return type(self)(**new_params)
-
-    def cpu(self):
-        return self.to(device=torch.device("cpu"))
-
-    def cuda(self):
-        return self.to(device=torch.device("cuda"))
-
-    # the following functions make sure **frame_data can be passed to functions
-    def __iter__(self):
-        for f in fields(self):
-            yield f.name
-
-    def __getitem__(self, key):
-        return getattr(self, key)
-
-    def __len__(self):
-        return len(fields(self))
-
-    @classmethod
-    def collate(cls, batch):
-        """
-        Given a list objects `batch` of class `cls`, collates them into a batched
-        representation suitable for processing with deep networks.
-        """
-
-        elem = batch[0]
-
-        if isinstance(elem, cls):
-            pointcloud_ids = [id(el.sequence_point_cloud) for el in batch]
-            id_to_idx = defaultdict(list)
-            for i, pc_id in enumerate(pointcloud_ids):
-                id_to_idx[pc_id].append(i)
-
-            sequence_point_cloud = []
-            sequence_point_cloud_idx = -np.ones((len(batch),))
-            for i, ind in enumerate(id_to_idx.values()):
-                sequence_point_cloud_idx[ind] = i
-                sequence_point_cloud.append(batch[ind[0]].sequence_point_cloud)
-            assert (sequence_point_cloud_idx >= 0).all()
-
-            override_fields = {
-                "sequence_point_cloud": sequence_point_cloud,
-                "sequence_point_cloud_idx": sequence_point_cloud_idx.tolist(),
-            }
-            # note that the pre-collate value of sequence_point_cloud_idx is unused
-
-            collated = {}
-            for f in fields(elem):
-                list_values = override_fields.get(
-                    f.name, [getattr(d, f.name) for d in batch]
-                )
-                collated[f.name] = (
-                    cls.collate(list_values)
-                    if all(list_value is not None for list_value in list_values)
-                    else None
-                )
-            return cls(**collated)
-
-        elif isinstance(elem, Pointclouds):
-            return join_pointclouds_as_batch(batch)
-
-        elif isinstance(elem, CamerasBase):
-            # TODO: don't store K; enforce working in NDC space
-            return join_cameras_as_batch(batch)
-        else:
-            return torch.utils.data._utils.collate.default_collate(batch)
-
-
-class _GenericWorkaround:
-    """
-    OmegaConf.structured has a weirdness when you try to apply
-    it to a dataclass whose first base class is a Generic which is not
-    Dict. The issue is with a function called get_dict_key_value_types
-    in omegaconf/_utils.py.
-    For example this fails:
-
-        @dataclass(eq=False)
-        class D(torch.utils.data.Dataset[int]):
-            a: int = 3
-
-        OmegaConf.structured(D)
-
-    We avoid the problem by adding this class as an extra base class.
-    """
-
-    pass
+from pytorch3d.implicitron.dataset.frame_data import FrameData
+from pytorch3d.implicitron.dataset.utils import GenericWorkaround
 
 
 @dataclass(eq=False)
-class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
+class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
     """
     Base class to describe a dataset to be used with Implicitron.
 
@@ -237,7 +49,7 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         raise NotImplementedError()
 
     def get_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int]
+        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
     ) -> List[Tuple[int, float]]:
         """
         If the sequences in the dataset are videos rather than
@@ -251,7 +63,9 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         frames.
 
         Args:
-            idx: frame index in self
+            idxs: frame index in self
+            subset_filter: If given, an index in idxs is ignored if the
+                corresponding frame is not in any of the named subsets.
 
         Returns:
             tuple of
@@ -291,7 +105,7 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         return dict(c2seq)
 
     def sequence_frames_in_order(
-        self, seq_name: str
+        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
     ) -> Iterator[Tuple[float, int, int]]:
         """Returns an iterator over the frame indices in a given sequence.
         We attempt to first sort by timestamp (if they are available),
@@ -308,7 +122,9 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
         """
         # pyre-ignore[16]
         seq_frame_indices = self._seq_to_idx[seq_name]
-        nos_timestamps = self.get_frame_numbers_and_timestamps(seq_frame_indices)
+        nos_timestamps = self.get_frame_numbers_and_timestamps(
+            seq_frame_indices, subset_filter
+        )
 
         yield from sorted(
             [
@@ -317,11 +133,13 @@ class DatasetBase(_GenericWorkaround, torch.utils.data.Dataset[FrameData]):
             ]
         )
 
-    def sequence_indices_in_order(self, seq_name: str) -> Iterator[int]:
+    def sequence_indices_in_order(
+        self, seq_name: str, subset_filter: Optional[Sequence[str]] = None
+    ) -> Iterator[int]:
         """Same as `sequence_frames_in_order` but returns the iterator over
         only dataset indices.
         """
-        for _, _, idx in self.sequence_frames_in_order(seq_name):
+        for _, _, idx in self.sequence_frames_in_order(seq_name, subset_filter):
             yield idx
 
     # frame_data_type is the actual type of frames returned by the dataset.

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
@@ -95,6 +97,7 @@ class DatasetMapProviderBase(ReplaceableBase):
 
     def get_all_train_cameras(self) -> Optional[CamerasBase]:
         """
+        DEPRECATED! The function will be removed in future versions.
         If the data is all for a single scene, returns a list
         of the known training cameras for that scene, which is
         used for evaluating the difficulty of the unknown

pytorch3d/implicitron/dataset/frame_data.py

@@ -0,0 +1,779 @@
+# 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 os
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass, field, fields
+from typing import (
+    Any,
+    ClassVar,
+    Generic,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
+
+import numpy as np
+import torch
+
+from pytorch3d.implicitron.dataset import types
+from pytorch3d.implicitron.dataset.utils import (
+    adjust_camera_to_bbox_crop_,
+    adjust_camera_to_image_scale_,
+    bbox_xyxy_to_xywh,
+    clamp_box_to_image_bounds_and_round,
+    crop_around_box,
+    GenericWorkaround,
+    get_bbox_from_mask,
+    get_clamp_bbox,
+    load_depth,
+    load_depth_mask,
+    load_image,
+    load_mask,
+    load_pointcloud,
+    rescale_bbox,
+    resize_image,
+    safe_as_tensor,
+)
+from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
+from pytorch3d.renderer.camera_utils import join_cameras_as_batch
+from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
+from pytorch3d.structures.pointclouds import join_pointclouds_as_batch, Pointclouds
+
+
+@dataclass
+class FrameData(Mapping[str, Any]):
+    """
+    A type of the elements returned by indexing the dataset object.
+    It can represent both individual frames and batches of thereof;
+    in this documentation, the sizes of tensors refer to single frames;
+    add the first batch dimension for the collation result.
+
+    Args:
+        frame_number: The number of the frame within its sequence.
+            0-based continuous integers.
+        sequence_name: The unique name of the frame's sequence.
+        sequence_category: The object category of the sequence.
+        frame_timestamp: The time elapsed since the start of a sequence in sec.
+        image_size_hw: The size of the original image in pixels; (height, width)
+            tensor of shape (2,). Note that it is optional, e.g. it can be `None`
+            if the frame annotation has no size ans image_rgb has not [yet] been
+            loaded. Image-less FrameData is valid but mutators like crop/resize
+            may fail if the original image size cannot be deduced.
+        effective_image_size_hw: The size of the image after mutations such as
+            crop/resize in pixels; (height, width). if the image has not been mutated,
+            it is equal to `image_size_hw`. Note that it is also optional, for the
+            same reason as `image_size_hw`.
+        image_path: The qualified path to the loaded image (with dataset_root).
+        image_rgb: A Tensor of shape `(3, H, W)` holding the RGB image
+            of the frame; elements are floats in [0, 1].
+        mask_crop: A binary mask of shape `(1, H, W)` denoting the valid image
+            regions. Regions can be invalid (mask_crop[i,j]=0) in case they
+            are a result of zero-padding of the image after cropping around
+            the object bounding box; elements are floats in {0.0, 1.0}.
+        depth_path: The qualified path to the frame's depth map.
+        depth_map: A float Tensor of shape `(1, H, W)` holding the depth map
+            of the frame; values correspond to distances from the camera;
+            use `depth_mask` and `mask_crop` to filter for valid pixels.
+        depth_mask: A binary mask of shape `(1, H, W)` denoting pixels of the
+            depth map that are valid for evaluation, they have been checked for
+            consistency across views; elements are floats in {0.0, 1.0}.
+        mask_path: A qualified path to the foreground probability mask.
+        fg_probability: A Tensor of `(1, H, W)` denoting the probability of the
+            pixels belonging to the captured object; elements are floats
+            in [0, 1].
+        bbox_xywh: The bounding box tightly enclosing the foreground object in the
+            format (x0, y0, width, height). The convention assumes that
+            `x0+width` and `y0+height` includes the boundary of the box.
+            I.e., to slice out the corresponding crop from an image tensor `I`
+            we execute `crop = I[..., y0:y0+height, x0:x0+width]`
+        crop_bbox_xywh: The bounding box denoting the boundaries of `image_rgb`
+            in the original image coordinates in the format (x0, y0, width, height).
+            The convention is the same as for `bbox_xywh`. `crop_bbox_xywh` differs
+            from `bbox_xywh` due to padding (which can happen e.g. due to
+            setting `JsonIndexDataset.box_crop_context > 0`)
+        camera: A PyTorch3D camera object corresponding the frame's viewpoint,
+            corrected for cropping if it happened.
+        camera_quality_score: The score proportional to the confidence of the
+            frame's camera estimation (the higher the more accurate).
+        point_cloud_quality_score: The score proportional to the accuracy of the
+            frame's sequence point cloud (the higher the more accurate).
+        sequence_point_cloud_path: The path to the sequence's point cloud.
+        sequence_point_cloud: A PyTorch3D Pointclouds object holding the
+            point cloud corresponding to the frame's sequence. When the object
+            represents a batch of frames, point clouds may be deduplicated;
+            see `sequence_point_cloud_idx`.
+        sequence_point_cloud_idx: Integer indices mapping frame indices to the
+            corresponding point clouds in `sequence_point_cloud`; to get the
+            corresponding point cloud to `image_rgb[i]`, use
+            `sequence_point_cloud[sequence_point_cloud_idx[i]]`.
+        frame_type: The type of the loaded frame specified in
+            `subset_lists_file`, if provided.
+        meta: A dict for storing additional frame information.
+    """
+
+    frame_number: Optional[torch.LongTensor]
+    sequence_name: Union[str, List[str]]
+    sequence_category: Union[str, List[str]]
+    frame_timestamp: Optional[torch.Tensor] = None
+    image_size_hw: Optional[torch.LongTensor] = None
+    effective_image_size_hw: Optional[torch.LongTensor] = None
+    image_path: Union[str, List[str], None] = None
+    image_rgb: Optional[torch.Tensor] = None
+    # masks out padding added due to cropping the square bit
+    mask_crop: Optional[torch.Tensor] = None
+    depth_path: Union[str, List[str], None] = None
+    depth_map: Optional[torch.Tensor] = None
+    depth_mask: Optional[torch.Tensor] = None
+    mask_path: Union[str, List[str], None] = None
+    fg_probability: Optional[torch.Tensor] = None
+    bbox_xywh: Optional[torch.Tensor] = None
+    crop_bbox_xywh: Optional[torch.Tensor] = None
+    camera: Optional[PerspectiveCameras] = None
+    camera_quality_score: Optional[torch.Tensor] = None
+    point_cloud_quality_score: Optional[torch.Tensor] = None
+    sequence_point_cloud_path: Union[str, List[str], None] = None
+    sequence_point_cloud: Optional[Pointclouds] = None
+    sequence_point_cloud_idx: Optional[torch.Tensor] = None
+    frame_type: Union[str, List[str], None] = None  # known | unseen
+    meta: dict = field(default_factory=lambda: {})
+
+    # NOTE that batching resets this attribute
+    _uncropped: bool = field(init=False, default=True)
+
+    def to(self, *args, **kwargs):
+        new_params = {}
+        for field_name in iter(self):
+            value = getattr(self, field_name)
+            if isinstance(value, (torch.Tensor, Pointclouds, CamerasBase)):
+                new_params[field_name] = value.to(*args, **kwargs)
+            else:
+                new_params[field_name] = value
+        frame_data = type(self)(**new_params)
+        frame_data._uncropped = self._uncropped
+        return frame_data
+
+    def cpu(self):
+        return self.to(device=torch.device("cpu"))
+
+    def cuda(self):
+        return self.to(device=torch.device("cuda"))
+
+    # the following functions make sure **frame_data can be passed to functions
+    def __iter__(self):
+        for f in fields(self):
+            if f.name.startswith("_"):
+                continue
+
+            yield f.name
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __len__(self):
+        return sum(1 for f in iter(self))
+
+    def crop_by_metadata_bbox_(
+        self,
+        box_crop_context: float,
+    ) -> None:
+        """Crops the frame data in-place by (possibly expanded) bounding box.
+        The bounding box is taken from the object state (usually taken from
+        the frame annotation or estimated from the foregroubnd mask).
+        If the expanded bounding box does not fit the image, it is clamped,
+        i.e. the image is *not* padded.
+
+        Args:
+            box_crop_context: rate of expansion for bbox; 0 means no expansion,
+
+        Raises:
+            ValueError: If the object does not contain a bounding box (usually when no
+                mask annotation is provided)
+            ValueError: If the frame data have been cropped or resized, thus the intrinsic
+                bounding box is not valid for the current image size.
+            ValueError: If the frame does not have an image size (usually a corner case
+                when no image has been loaded)
+        """
+        if self.bbox_xywh is None:
+            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(
+                "Trying to apply the metadata bounding box to already cropped "
+                "or resized image; coordinates have changed."
+            )
+
+        self._crop_by_bbox_(
+            box_crop_context,
+            self.bbox_xywh,
+        )
+
+    def crop_by_given_bbox_(
+        self,
+        box_crop_context: float,
+        bbox_xywh: torch.Tensor,
+    ) -> None:
+        """Crops the frame data in-place by (possibly expanded) bounding box.
+        If the expanded bounding box does not fit the image, it is clamped,
+        i.e. the image is *not* padded.
+
+        Args:
+            box_crop_context: rate of expansion for bbox; 0 means no expansion,
+            bbox_xywh: bounding box in [x0, y0, width, height] format. If float
+                tensor, values are floored (after converting to [x0, y0, x1, y1]).
+
+        Raises:
+            ValueError: If the frame does not have an image size (usually a corner case
+                when no image has been loaded)
+        """
+        self._crop_by_bbox_(
+            box_crop_context,
+            bbox_xywh,
+        )
+
+    def _crop_by_bbox_(
+        self,
+        box_crop_context: float,
+        bbox_xywh: torch.Tensor,
+    ) -> None:
+        """Crops the frame data in-place by (possibly expanded) bounding box.
+        If the expanded bounding box does not fit the image, it is clamped,
+        i.e. the image is *not* padded.
+
+        Args:
+            box_crop_context: rate of expansion for bbox; 0 means no expansion,
+            bbox_xywh: bounding box in [x0, y0, width, height] format. If float
+                tensor, values are floored (after converting to [x0, y0, x1, y1]).
+
+        Raises:
+            ValueError: If the frame does not have an image size (usually a corner case
+                when no image has been loaded)
+        """
+        effective_image_size_hw = self.effective_image_size_hw
+        if effective_image_size_hw is None:
+            raise ValueError("Calling crop on image-less FrameData")
+
+        bbox_xyxy = get_clamp_bbox(
+            bbox_xywh,
+            image_path=self.image_path,  # pyre-ignore
+            box_crop_context=box_crop_context,
+        )
+        clamp_bbox_xyxy = clamp_box_to_image_bounds_and_round(
+            bbox_xyxy,
+            image_size_hw=tuple(self.effective_image_size_hw),  # pyre-ignore
+        )
+        crop_bbox_xywh = bbox_xyxy_to_xywh(clamp_bbox_xyxy)
+
+        if self.fg_probability is not None:
+            self.fg_probability = crop_around_box(
+                self.fg_probability,
+                clamp_bbox_xyxy,
+                self.mask_path,  # pyre-ignore
+            )
+        if self.image_rgb is not None:
+            self.image_rgb = crop_around_box(
+                self.image_rgb,
+                clamp_bbox_xyxy,
+                self.image_path,  # pyre-ignore
+            )
+
+        depth_map = self.depth_map
+        if depth_map is not None:
+            clamp_bbox_xyxy_depth = rescale_bbox(
+                clamp_bbox_xyxy, tuple(depth_map.shape[-2:]), effective_image_size_hw
+            ).long()
+            self.depth_map = crop_around_box(
+                depth_map,
+                clamp_bbox_xyxy_depth,
+                self.depth_path,  # pyre-ignore
+            )
+
+        depth_mask = self.depth_mask
+        if depth_mask is not None:
+            clamp_bbox_xyxy_depth = rescale_bbox(
+                clamp_bbox_xyxy, tuple(depth_mask.shape[-2:]), effective_image_size_hw
+            ).long()
+            self.depth_mask = crop_around_box(
+                depth_mask,
+                clamp_bbox_xyxy_depth,
+                self.mask_path,  # pyre-ignore
+            )
+
+        # changing principal_point according to bbox_crop
+        if self.camera is not None:
+            adjust_camera_to_bbox_crop_(
+                camera=self.camera,
+                image_size_wh=effective_image_size_hw.flip(dims=[-1]),
+                clamp_bbox_xywh=crop_bbox_xywh,
+            )
+
+        # pyre-ignore
+        self.effective_image_size_hw = crop_bbox_xywh[..., 2:].flip(dims=[-1])
+        self._uncropped = False
+
+    def resize_frame_(self, new_size_hw: torch.LongTensor) -> None:
+        """Resizes frame data in-place according to given dimensions.
+
+        Args:
+            new_size_hw: target image size [height, width], a LongTensor of shape (2,)
+
+        Raises:
+            ValueError: If the frame does not have an image size (usually a corner case
+                when no image has been loaded)
+        """
+
+        effective_image_size_hw = self.effective_image_size_hw
+        if effective_image_size_hw is None:
+            raise ValueError("Calling resize on image-less FrameData")
+
+        image_height, image_width = new_size_hw.tolist()
+
+        if self.fg_probability is not None:
+            self.fg_probability, _, _ = resize_image(
+                self.fg_probability,
+                image_height=image_height,
+                image_width=image_width,
+                mode="nearest",
+            )
+
+        if self.image_rgb is not None:
+            self.image_rgb, _, self.mask_crop = resize_image(
+                self.image_rgb, image_height=image_height, image_width=image_width
+            )
+
+        if self.depth_map is not None:
+            self.depth_map, _, _ = resize_image(
+                self.depth_map,
+                image_height=image_height,
+                image_width=image_width,
+                mode="nearest",
+            )
+
+        if self.depth_mask is not None:
+            self.depth_mask, _, _ = resize_image(
+                self.depth_mask,
+                image_height=image_height,
+                image_width=image_width,
+                mode="nearest",
+            )
+
+        if self.camera is not None:
+            if self.image_size_hw is None:
+                raise ValueError(
+                    "image_size_hw has to be defined for resizing FrameData with cameras."
+                )
+            adjust_camera_to_image_scale_(
+                camera=self.camera,
+                original_size_wh=effective_image_size_hw.flip(dims=[-1]),
+                new_size_wh=new_size_hw.flip(dims=[-1]),  # pyre-ignore
+            )
+
+        self.effective_image_size_hw = new_size_hw
+        self._uncropped = False
+
+    @classmethod
+    def collate(cls, batch):
+        """
+        Given a list objects `batch` of class `cls`, collates them into a batched
+        representation suitable for processing with deep networks.
+        """
+
+        elem = batch[0]
+
+        if isinstance(elem, cls):
+            pointcloud_ids = [id(el.sequence_point_cloud) for el in batch]
+            id_to_idx = defaultdict(list)
+            for i, pc_id in enumerate(pointcloud_ids):
+                id_to_idx[pc_id].append(i)
+
+            sequence_point_cloud = []
+            sequence_point_cloud_idx = -np.ones((len(batch),))
+            for i, ind in enumerate(id_to_idx.values()):
+                sequence_point_cloud_idx[ind] = i
+                sequence_point_cloud.append(batch[ind[0]].sequence_point_cloud)
+            assert (sequence_point_cloud_idx >= 0).all()
+
+            override_fields = {
+                "sequence_point_cloud": sequence_point_cloud,
+                "sequence_point_cloud_idx": sequence_point_cloud_idx.tolist(),
+            }
+            # note that the pre-collate value of sequence_point_cloud_idx is unused
+
+            collated = {}
+            for f in fields(elem):
+                if not f.init:
+                    continue
+
+                list_values = override_fields.get(
+                    f.name, [getattr(d, f.name) for d in batch]
+                )
+                collated[f.name] = (
+                    cls.collate(list_values)
+                    if all(list_value is not None for list_value in list_values)
+                    else None
+                )
+            return cls(**collated)
+
+        elif isinstance(elem, Pointclouds):
+            return join_pointclouds_as_batch(batch)
+
+        elif isinstance(elem, CamerasBase):
+            # TODO: don't store K; enforce working in NDC space
+            return join_cameras_as_batch(batch)
+        else:
+            return torch.utils.data._utils.collate.default_collate(batch)
+
+
+FrameDataSubtype = TypeVar("FrameDataSubtype", bound=FrameData)
+
+
+class FrameDataBuilderBase(ReplaceableBase, Generic[FrameDataSubtype], ABC):
+    """A base class for FrameDataBuilders that build a FrameData object, load and
+    process the binary data (crop and resize). Implementations should parametrize
+    the class with a subtype of FrameData and set frame_data_type class variable to
+    that type. They have to also implement `build` method.
+    """
+
+    # To be initialised to FrameDataSubtype
+    frame_data_type: ClassVar[Type[FrameDataSubtype]]
+
+    @abstractmethod
+    def build(
+        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.
+        """
+        raise NotImplementedError()
+
+
+class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
+    """
+    A class to build a FrameData object, load and process the binary data (crop and
+    resize). This is an abstract class for extending to build FrameData subtypes. Most
+    users need to use concrete `FrameDataBuilder` class instead.
+    Beware that modifications of frame data are done in-place.
+
+    Args:
+        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
+            depth values used for evaluation (the points consistent across views).
+        load_masks: Enable loading frame foreground masks.
+        load_point_clouds: Enable loading sequence-level point clouds.
+        max_points: Cap on the number of loaded points in the point cloud;
+            if reached, they are randomly sampled without replacement.
+        mask_images: Whether to mask the images with the loaded foreground masks;
+            0 value is used for background.
+        mask_depths: Whether to mask the depth maps with the loaded foreground
+            masks; 0 value is used for background.
+        image_height: The height of the returned images, masks, and depth maps;
+            aspect ratio is preserved during cropping/resizing.
+        image_width: The width of the returned images, masks, and depth maps;
+            aspect ratio is preserved during cropping/resizing.
+        box_crop: Enable cropping of the image around the bounding box inferred
+            from the foreground region of the loaded segmentation mask; masks
+            and depth maps are cropped accordingly; cameras are corrected.
+        box_crop_mask_thr: The threshold used to separate pixels into foreground
+            and background based on the foreground_probability mask; if no value
+            is greater than this threshold, the loader lowers it and repeats.
+        box_crop_context: The amount of additional padding added to each
+            dimension of the cropping bounding box, relative to box size.
+        path_manager: Optionally a PathManager for interpreting paths in a special way.
+    """
+
+    dataset_root: Optional[str] = None
+    load_images: bool = True
+    load_depths: bool = True
+    load_depth_masks: bool = True
+    load_masks: bool = True
+    load_point_clouds: bool = False
+    max_points: int = 0
+    mask_images: bool = False
+    mask_depths: bool = False
+    image_height: Optional[int] = 800
+    image_width: Optional[int] = 800
+    box_crop: bool = True
+    box_crop_mask_thr: float = 0.4
+    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:
+            * if box_crop is set, the image/mask/depth are cropped with the bounding
+                box provided or estimated from MaskAnnotation,
+            * if image_height/image_width are set, the image/mask/depth are resized to
+                fit that resolution. Note that the aspect ratio is preserved, and the
+                (possibly cropped) image is pasted into the top-left corner. In the
+                resulting frame_data, mask_crop field corresponds to the mask of the
+                pasted image.
+
+        Args:
+            frame_annotation: frame annotation
+            sequence_annotation: sequence annotation
+            load_blobs: if the function should attempt loading the image, depth map
+                and mask, and foreground mask
+
+        Returns:
+            The constructed FrameData object.
+        """
+
+        point_cloud = sequence_annotation.point_cloud
+
+        frame_data = self.frame_data_type(
+            frame_number=safe_as_tensor(frame_annotation.frame_number, torch.long),
+            frame_timestamp=safe_as_tensor(
+                frame_annotation.frame_timestamp, torch.float
+            ),
+            sequence_name=frame_annotation.sequence_name,
+            sequence_category=sequence_annotation.category,
+            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
+            ),
+        )
+
+        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:
+                if image_path is None:
+                    raise ValueError("Image path is required to load images.")
+
+                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, 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)
+            frame_data.sequence_point_cloud = load_pointcloud(
+                self._local_path(pcl_path), max_points=self.max_points
+            )
+            frame_data.sequence_point_cloud_path = pcl_path
+
+        if frame_annotation.viewpoint is not None:
+            frame_data.camera = self._get_pytorch3d_camera(frame_annotation)
+
+        if self.box_crop:
+            frame_data.crop_by_metadata_bbox_(self.box_crop_context)
+
+        if self.image_height is not None and self.image_width is not None:
+            new_size = (self.image_height, self.image_width)
+            frame_data.resize_frame_(
+                new_size_hw=torch.tensor(new_size, dtype=torch.long),  # pyre-ignore
+            )
+
+        return frame_data
+
+    def _load_fg_probability(
+        self, entry: types.FrameAnnotation
+    ) -> 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))
+        if fg_probability.shape[-2:] != entry.image.size:
+            raise ValueError(
+                f"bad mask size: {fg_probability.shape[-2:]} vs {entry.image.size}!"
+            )
+
+        return fg_probability, full_path
+
+    def _postprocess_image(
+        self,
+        image_np: np.ndarray,
+        image_size: Tuple[int, int],
+        fg_probability: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        image_rgb = safe_as_tensor(image_np, torch.float)
+
+        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
+
+    def _load_mask_depth(
+        self,
+        entry: types.FrameAnnotation,
+        fg_mask: Optional[np.ndarray],
+    ) -> Tuple[torch.Tensor, str, torch.Tensor]:
+        entry_depth = entry.depth
+        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_mask is not None
+            depth_map *= fg_mask
+
+        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 = (depth_map > 0.0).astype(np.float32)
+
+        return torch.tensor(depth_map), path, torch.tensor(depth_mask)
+
+    def _get_pytorch3d_camera(
+        self,
+        entry: types.FrameAnnotation,
+    ) -> PerspectiveCameras:
+        entry_viewpoint = entry.viewpoint
+        assert entry_viewpoint is not None
+        # principal point and focal length
+        principal_point = torch.tensor(
+            entry_viewpoint.principal_point, dtype=torch.float
+        )
+        focal_length = torch.tensor(entry_viewpoint.focal_length, dtype=torch.float)
+
+        format = entry_viewpoint.intrinsics_format
+        if entry_viewpoint.intrinsics_format == "ndc_norm_image_bounds":
+            # legacy PyTorch3D NDC format
+            # convert to pixels unequally and convert to ndc equally
+            image_size_as_list = list(reversed(entry.image.size))
+            image_size_wh = torch.tensor(image_size_as_list, dtype=torch.float)
+            per_axis_scale = image_size_wh / image_size_wh.min()
+            focal_length = focal_length * per_axis_scale
+            principal_point = principal_point * per_axis_scale
+        elif entry_viewpoint.intrinsics_format != "ndc_isotropic":
+            raise ValueError(f"Unknown intrinsics format: {format}")
+
+        return PerspectiveCameras(
+            focal_length=focal_length[None],
+            principal_point=principal_point[None],
+            R=torch.tensor(entry_viewpoint.R, dtype=torch.float)[None],
+            T=torch.tensor(entry_viewpoint.T, dtype=torch.float)[None],
+        )
+
+    def _fix_point_cloud_path(self, path: str) -> str:
+        """
+        Fix up a point cloud path from the dataset.
+        Some files in Co3Dv2 have an accidental absolute path stored.
+        """
+        unwanted_prefix = (
+            "/large_experiments/p3/replay/datasets/co3d/co3d45k_220512/export_v23/"
+        )
+        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:
+        if self.path_manager is None:
+            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]):
+    """
+    A concrete class to build a FrameData object, load and process the binary data (crop
+    and resize). Beware that modifications of frame data are done in-place. Please see
+    the documentation for `GenericFrameDataBuilder` for the description of parameters
+    and methods.
+    """
+
+    frame_data_type: ClassVar[Type[FrameData]] = 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
@@ -15,7 +17,6 @@ import random
 import warnings
 from collections import defaultdict
 from itertools import islice
-from pathlib import Path
 from typing import (
     Any,
     ClassVar,
@@ -30,19 +31,15 @@ from typing import (
     Union,
 )
 
-import numpy as np
-import torch
-from PIL import Image
+from pytorch3d.implicitron.dataset import types
+from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
+from pytorch3d.implicitron.dataset.frame_data import FrameData, FrameDataBuilder
+from pytorch3d.implicitron.dataset.utils import is_known_frame_scalar
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
-from pytorch3d.io import IO
 from pytorch3d.renderer.camera_utils import join_cameras_as_batch
-from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
-from pytorch3d.structures.pointclouds import Pointclouds
-from tqdm import tqdm
+from pytorch3d.renderer.cameras import CamerasBase
 
-from . import types
-from .dataset_base import DatasetBase, FrameData
-from .utils import is_known_frame_scalar
+from tqdm import tqdm
 
 
 logger = logging.getLogger(__name__)
@@ -65,7 +62,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
     A dataset with annotations in json files like the Common Objects in 3D
     (CO3D) dataset.
 
-    Args:
+    Metadata-related args::
         frame_annotations_file: A zipped json file containing metadata of the
             frames in the dataset, serialized List[types.FrameAnnotation].
         sequence_annotations_file: A zipped json file containing metadata of the
@@ -83,6 +80,24 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         pick_sequence: A list of sequence names to restrict the dataset to.
         exclude_sequence: A list of the names of the sequences to exclude.
         limit_category_to: Restrict the dataset to the given list of categories.
+        remove_empty_masks: Removes the frames with no active foreground pixels
+            in the segmentation mask after thresholding (see box_crop_mask_thr).
+        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 before other frame-level filters.
+        seed: The seed of the random generator sampling #n_frames_per_sequence
+            random frames per sequence.
+        sort_frames: Enable frame annotations sorting to group frames from the
+            same sequences together and order them by timestamps
+        eval_batches: A list of batches that form the evaluation set;
+            list of batch-sized lists of indices corresponding to __getitem__
+            of this class, thus it can be used directly as a batch sampler.
+        eval_batch_index:
+            ( Optional[List[List[Union[Tuple[str, int, str], Tuple[str, int]]]] )
+            A list of batches of frames described as (sequence_name, frame_idx)
+            that can form the evaluation set, `eval_batches` will be set from this.
+
+    Blob-loading parameters:
         dataset_root: The root folder of the dataset; all the paths in jsons are
             specified relative to this root (but not json paths themselves).
         load_images: Enable loading the frame RGB data.
@@ -109,28 +124,11 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             is greater than this threshold, the loader lowers it and repeats.
         box_crop_context: The amount of additional padding added to each
             dimension of the cropping bounding box, relative to box size.
-        remove_empty_masks: Removes the frames with no active foreground pixels
-            in the segmentation mask after thresholding (see box_crop_mask_thr).
-        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 before other frame-level filters.
-        seed: The seed of the random generator sampling #n_frames_per_sequence
-            random frames per sequence.
-        sort_frames: Enable frame annotations sorting to group frames from the
-            same sequences together and order them by timestamps
-        eval_batches: A list of batches that form the evaluation set;
-            list of batch-sized lists of indices corresponding to __getitem__
-            of this class, thus it can be used directly as a batch sampler.
-        eval_batch_index:
-            ( Optional[List[List[Union[Tuple[str, int, str], Tuple[str, int]]]] )
-            A list of batches of frames described as (sequence_name, frame_idx)
-            that can form the evaluation set, `eval_batches` will be set from this.
-
     """
 
-    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 = ""
@@ -162,12 +160,14 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
     sort_frames: bool = False
     eval_batches: Any = None
     eval_batch_index: Any = None
+    # initialised in __post_init__
+    # commented because of OmegaConf (for tests to pass)
+    # _frame_data_builder: FrameDataBuilder = field(init=False)
     # frame_annots: List[FrameAnnotsEntry] = field(init=False)
     # seq_annots: Dict[str, types.SequenceAnnotation] = field(init=False)
+    # _seq_to_idx: Dict[str, List[int]] = field(init=False)
 
     def __post_init__(self) -> None:
-        # pyre-fixme[16]: `JsonIndexDataset` has no attribute `subset_to_image_path`.
-        self.subset_to_image_path = None
         self._load_frames()
         self._load_sequences()
         if self.sort_frames:
@@ -175,9 +175,28 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         self._load_subset_lists()
         self._filter_db()  # also computes sequence indices
         self._extract_and_set_eval_batches()
+
+        # pyre-ignore
+        self._frame_data_builder = FrameDataBuilder(
+            dataset_root=self.dataset_root,
+            load_images=self.load_images,
+            load_depths=self.load_depths,
+            load_depth_masks=self.load_depth_masks,
+            load_masks=self.load_masks,
+            load_point_clouds=self.load_point_clouds,
+            max_points=self.max_points,
+            mask_images=self.mask_images,
+            mask_depths=self.mask_depths,
+            image_height=self.image_height,
+            image_width=self.image_width,
+            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))
 
-    def _extract_and_set_eval_batches(self):
+    def _extract_and_set_eval_batches(self) -> None:
         """
         Sets eval_batches based on input eval_batch_index.
         """
@@ -207,13 +226,13 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             # https://gist.github.com/treyhunner/f35292e676efa0be1728
             functools.reduce(
                 lambda a, b: {**a, **b},
-                [d.seq_annots for d in other_datasets],  # pyre-ignore[16]
+                # pyre-ignore[16]
+                [d.seq_annots for d in other_datasets],
             )
         )
         all_eval_batches = [
             self.eval_batches,
-            # pyre-ignore
-            *[d.eval_batches for d in other_datasets],
+            *[d.eval_batches for d in other_datasets],  # pyre-ignore[16]
         ]
         if not (
             all(ba is None for ba in all_eval_batches)
@@ -251,7 +270,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         allow_missing_indices: bool = False,
         remove_missing_indices: bool = False,
         suppress_missing_index_warning: bool = True,
-    ) -> List[List[Union[Optional[int], int]]]:
+    ) -> Union[List[List[Optional[int]]], List[List[int]]]:
         """
         Obtain indices into the dataset object given a list of frame ids.
 
@@ -323,9 +342,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             valid_dataset_idx = [
                 [b for b in batch if b is not None] for batch in dataset_idx
             ]
-            return [  # pyre-ignore[7]
-                batch for batch in valid_dataset_idx if len(batch) > 0
-            ]
+            return [batch for batch in valid_dataset_idx if len(batch) > 0]
 
         return dataset_idx
 
@@ -417,255 +434,18 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
             raise IndexError(f"index {index} out of range {len(self.frame_annots)}")
 
         entry = self.frame_annots[index]["frame_annotation"]
-        # pyre-ignore[16]
-        point_cloud = self.seq_annots[entry.sequence_name].point_cloud
-        frame_data = FrameData(
-            frame_number=_safe_as_tensor(entry.frame_number, torch.long),
-            frame_timestamp=_safe_as_tensor(entry.frame_timestamp, torch.float),
-            sequence_name=entry.sequence_name,
-            sequence_category=self.seq_annots[entry.sequence_name].category,
-            camera_quality_score=_safe_as_tensor(
-                self.seq_annots[entry.sequence_name].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,
-        )
 
-        # The rest of the fields are optional
+        # pyre-ignore
+        frame_data = self._frame_data_builder.build(
+            entry,
+            # pyre-ignore
+            self.seq_annots[entry.sequence_name],
+        )
+        # Optional field
         frame_data.frame_type = self._get_frame_type(self.frame_annots[index])
 
-        (
-            frame_data.fg_probability,
-            frame_data.mask_path,
-            frame_data.bbox_xywh,
-            clamp_bbox_xyxy,
-            frame_data.crop_bbox_xywh,
-        ) = self._load_crop_fg_probability(entry)
-
-        scale = 1.0
-        if self.load_images and entry.image is not None:
-            # original image size
-            frame_data.image_size_hw = _safe_as_tensor(entry.image.size, torch.long)
-
-            (
-                frame_data.image_rgb,
-                frame_data.image_path,
-                frame_data.mask_crop,
-                scale,
-            ) = self._load_crop_images(
-                entry, frame_data.fg_probability, clamp_bbox_xyxy
-            )
-
-        if self.load_depths and entry.depth is not None:
-            (
-                frame_data.depth_map,
-                frame_data.depth_path,
-                frame_data.depth_mask,
-            ) = self._load_mask_depth(entry, clamp_bbox_xyxy, frame_data.fg_probability)
-
-        if entry.viewpoint is not None:
-            frame_data.camera = self._get_pytorch3d_camera(
-                entry,
-                scale,
-                clamp_bbox_xyxy,
-            )
-
-        if self.load_point_clouds and point_cloud is not None:
-            pcl_path = self._fix_point_cloud_path(point_cloud.path)
-            frame_data.sequence_point_cloud = _load_pointcloud(
-                self._local_path(pcl_path), max_points=self.max_points
-            )
-            frame_data.sequence_point_cloud_path = pcl_path
-
         return frame_data
 
-    def _fix_point_cloud_path(self, path: str) -> str:
-        """
-        Fix up a point cloud path from the dataset.
-        Some files in Co3Dv2 have an accidental absolute path stored.
-        """
-        unwanted_prefix = (
-            "/large_experiments/p3/replay/datasets/co3d/co3d45k_220512/export_v23/"
-        )
-        if path.startswith(unwanted_prefix):
-            path = path[len(unwanted_prefix) :]
-        return os.path.join(self.dataset_root, path)
-
-    def _load_crop_fg_probability(
-        self, entry: types.FrameAnnotation
-    ) -> Tuple[
-        Optional[torch.Tensor],
-        Optional[str],
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-    ]:
-        fg_probability = None
-        full_path = None
-        bbox_xywh = None
-        clamp_bbox_xyxy = None
-        crop_box_xywh = None
-
-        if (self.load_masks or self.box_crop) and entry.mask is not None:
-            full_path = os.path.join(self.dataset_root, entry.mask.path)
-            mask = _load_mask(self._local_path(full_path))
-
-            if mask.shape[-2:] != entry.image.size:
-                raise ValueError(
-                    f"bad mask size: {mask.shape[-2:]} vs {entry.image.size}!"
-                )
-
-            bbox_xywh = torch.tensor(_get_bbox_from_mask(mask, self.box_crop_mask_thr))
-
-            if self.box_crop:
-                clamp_bbox_xyxy = _clamp_box_to_image_bounds_and_round(
-                    _get_clamp_bbox(
-                        bbox_xywh,
-                        image_path=entry.image.path,
-                        box_crop_context=self.box_crop_context,
-                    ),
-                    image_size_hw=tuple(mask.shape[-2:]),
-                )
-                crop_box_xywh = _bbox_xyxy_to_xywh(clamp_bbox_xyxy)
-
-                mask = _crop_around_box(mask, clamp_bbox_xyxy, full_path)
-
-            fg_probability, _, _ = self._resize_image(mask, mode="nearest")
-
-        return fg_probability, full_path, bbox_xywh, clamp_bbox_xyxy, crop_box_xywh
-
-    def _load_crop_images(
-        self,
-        entry: types.FrameAnnotation,
-        fg_probability: Optional[torch.Tensor],
-        clamp_bbox_xyxy: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, str, torch.Tensor, float]:
-        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))
-
-        if image_rgb.shape[-2:] != entry.image.size:
-            raise ValueError(
-                f"bad image size: {image_rgb.shape[-2:]} vs {entry.image.size}!"
-            )
-
-        if self.box_crop:
-            assert clamp_bbox_xyxy is not None
-            image_rgb = _crop_around_box(image_rgb, clamp_bbox_xyxy, path)
-
-        image_rgb, scale, mask_crop = self._resize_image(image_rgb)
-
-        if self.mask_images:
-            assert fg_probability is not None
-            image_rgb *= fg_probability
-
-        return image_rgb, path, mask_crop, scale
-
-    def _load_mask_depth(
-        self,
-        entry: types.FrameAnnotation,
-        clamp_bbox_xyxy: Optional[torch.Tensor],
-        fg_probability: Optional[torch.Tensor],
-    ) -> 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)
-        depth_map = _load_depth(self._local_path(path), entry_depth.scale_adjustment)
-
-        if self.box_crop:
-            assert clamp_bbox_xyxy is not None
-            depth_bbox_xyxy = _rescale_bbox(
-                clamp_bbox_xyxy, entry.image.size, depth_map.shape[-2:]
-            )
-            depth_map = _crop_around_box(depth_map, depth_bbox_xyxy, path)
-
-        depth_map, _, _ = self._resize_image(depth_map, mode="nearest")
-
-        if self.mask_depths:
-            assert fg_probability is not None
-            depth_map *= fg_probability
-
-        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)
-            depth_mask = _load_depth_mask(self._local_path(mask_path))
-
-            if self.box_crop:
-                assert clamp_bbox_xyxy is not None
-                depth_mask_bbox_xyxy = _rescale_bbox(
-                    clamp_bbox_xyxy, entry.image.size, depth_mask.shape[-2:]
-                )
-                depth_mask = _crop_around_box(
-                    depth_mask, depth_mask_bbox_xyxy, mask_path
-                )
-
-            depth_mask, _, _ = self._resize_image(depth_mask, mode="nearest")
-        else:
-            depth_mask = torch.ones_like(depth_map)
-
-        return depth_map, path, depth_mask
-
-    def _get_pytorch3d_camera(
-        self,
-        entry: types.FrameAnnotation,
-        scale: float,
-        clamp_bbox_xyxy: Optional[torch.Tensor],
-    ) -> PerspectiveCameras:
-        entry_viewpoint = entry.viewpoint
-        assert entry_viewpoint is not None
-        # principal point and focal length
-        principal_point = torch.tensor(
-            entry_viewpoint.principal_point, dtype=torch.float
-        )
-        focal_length = torch.tensor(entry_viewpoint.focal_length, dtype=torch.float)
-
-        half_image_size_wh_orig = (
-            torch.tensor(list(reversed(entry.image.size)), dtype=torch.float) / 2.0
-        )
-
-        # first, we convert from the dataset's NDC convention to pixels
-        format = entry_viewpoint.intrinsics_format
-        if format.lower() == "ndc_norm_image_bounds":
-            # this is e.g. currently used in CO3D for storing intrinsics
-            rescale = half_image_size_wh_orig
-        elif format.lower() == "ndc_isotropic":
-            rescale = half_image_size_wh_orig.min()
-        else:
-            raise ValueError(f"Unknown intrinsics format: {format}")
-
-        # principal point and focal length in pixels
-        principal_point_px = half_image_size_wh_orig - principal_point * rescale
-        focal_length_px = focal_length * rescale
-        if self.box_crop:
-            assert clamp_bbox_xyxy is not None
-            principal_point_px -= clamp_bbox_xyxy[:2]
-
-        # now, convert from pixels to PyTorch3D v0.5+ NDC convention
-        if self.image_height is None or self.image_width is None:
-            out_size = list(reversed(entry.image.size))
-        else:
-            out_size = [self.image_width, self.image_height]
-
-        half_image_size_output = torch.tensor(out_size, dtype=torch.float) / 2.0
-        half_min_image_size_output = half_image_size_output.min()
-
-        # rescaled principal point and focal length in ndc
-        principal_point = (
-            half_image_size_output - principal_point_px * scale
-        ) / half_min_image_size_output
-        focal_length = focal_length_px * scale / half_min_image_size_output
-
-        return PerspectiveCameras(
-            focal_length=focal_length[None],
-            principal_point=principal_point[None],
-            R=torch.tensor(entry_viewpoint.R, dtype=torch.float)[None],
-            T=torch.tensor(entry_viewpoint.T, dtype=torch.float)[None],
-        )
-
     def _load_frames(self) -> None:
         logger.info(f"Loading Co3D frames from {self.frame_annotations_file}.")
         local_file = self._local_path(self.frame_annotations_file)
@@ -853,46 +633,23 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
         # pyre-ignore[16]
         self._seq_to_idx = seq_to_idx
 
-    def _resize_image(
-        self, image, mode="bilinear"
-    ) -> Tuple[torch.Tensor, float, torch.Tensor]:
-        image_height, image_width = self.image_height, self.image_width
-        if image_height is None or image_width is None:
-            # skip the resizing
-            imre_ = torch.from_numpy(image)
-            return imre_, 1.0, torch.ones_like(imre_[:1])
-        # takes numpy array, returns pytorch tensor
-        minscale = min(
-            image_height / image.shape[-2],
-            image_width / image.shape[-1],
-        )
-        imre = torch.nn.functional.interpolate(
-            torch.from_numpy(image)[None],
-            scale_factor=minscale,
-            mode=mode,
-            align_corners=False if mode == "bilinear" else None,
-            recompute_scale_factor=True,
-        )[0]
-        # pyre-fixme[19]: Expected 1 positional argument.
-        imre_ = torch.zeros(image.shape[0], self.image_height, self.image_width)
-        imre_[:, 0 : imre.shape[1], 0 : imre.shape[2]] = imre
-        # pyre-fixme[6]: For 2nd param expected `int` but got `Optional[int]`.
-        # pyre-fixme[6]: For 3rd param expected `int` but got `Optional[int]`.
-        mask = torch.zeros(1, self.image_height, self.image_width)
-        mask[:, 0 : imre.shape[1], 0 : imre.shape[2]] = 1.0
-        return imre_, minscale, mask
-
     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_frame_numbers_and_timestamps(
-        self, idxs: Sequence[int]
+        self, idxs: Sequence[int], subset_filter: Optional[Sequence[str]] = None
     ) -> List[Tuple[int, float]]:
         out: List[Tuple[int, float]] = []
         for idx in idxs:
-            # pyre-ignore[16]
+            if (
+                subset_filter is not None
+                # pyre-fixme[16]: `JsonIndexDataset` has no attribute `frame_annots`.
+                and self.frame_annots[idx]["subset"] not in subset_filter
+            ):
+                continue
+
             frame_annotation = self.frame_annots[idx]["frame_annotation"]
             out.append(
                 (frame_annotation.frame_number, frame_annotation.frame_timestamp)
@@ -912,169 +669,3 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
 
 def _seq_name_to_seed(seq_name) -> int:
     return int(hashlib.sha1(seq_name.encode("utf-8")).hexdigest(), 16)
-
-
-def _load_image(path) -> np.ndarray:
-    with Image.open(path) as pil_im:
-        im = np.array(pil_im.convert("RGB"))
-    im = im.transpose((2, 0, 1))
-    im = im.astype(np.float32) / 255.0
-    return im
-
-
-def _load_16big_png_depth(depth_png) -> np.ndarray:
-    with Image.open(depth_png) as depth_pil:
-        # the image is stored with 16-bit depth but PIL reads it as I (32 bit).
-        # we cast it to uint16, then reinterpret as float16, then cast to float32
-        depth = (
-            np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
-            .astype(np.float32)
-            .reshape((depth_pil.size[1], depth_pil.size[0]))
-        )
-    return depth
-
-
-def _load_1bit_png_mask(file: str) -> np.ndarray:
-    with Image.open(file) as pil_im:
-        mask = (np.array(pil_im.convert("L")) > 0.0).astype(np.float32)
-    return mask
-
-
-def _load_depth_mask(path: str) -> np.ndarray:
-    if not path.lower().endswith(".png"):
-        raise ValueError('unsupported depth mask file name "%s"' % path)
-    m = _load_1bit_png_mask(path)
-    return m[None]  # fake feature channel
-
-
-def _load_depth(path, scale_adjustment) -> np.ndarray:
-    if not path.lower().endswith(".png"):
-        raise ValueError('unsupported depth file name "%s"' % path)
-
-    d = _load_16big_png_depth(path) * scale_adjustment
-    d[~np.isfinite(d)] = 0.0
-    return d[None]  # fake feature channel
-
-
-def _load_mask(path) -> np.ndarray:
-    with Image.open(path) as pil_im:
-        mask = np.array(pil_im)
-    mask = mask.astype(np.float32) / 255.0
-    return mask[None]  # fake feature channel
-
-
-def _get_1d_bounds(arr) -> Tuple[int, int]:
-    nz = np.flatnonzero(arr)
-    return nz[0], nz[-1] + 1
-
-
-def _get_bbox_from_mask(
-    mask, thr, decrease_quant: float = 0.05
-) -> Tuple[int, int, int, int]:
-    # bbox in xywh
-    masks_for_box = np.zeros_like(mask)
-    while masks_for_box.sum() <= 1.0:
-        masks_for_box = (mask > thr).astype(np.float32)
-        thr -= decrease_quant
-    if thr <= 0.0:
-        warnings.warn(f"Empty masks_for_bbox (thr={thr}) => using full image.")
-
-    x0, x1 = _get_1d_bounds(masks_for_box.sum(axis=-2))
-    y0, y1 = _get_1d_bounds(masks_for_box.sum(axis=-1))
-
-    return x0, y0, x1 - x0, y1 - y0
-
-
-def _get_clamp_bbox(
-    bbox: torch.Tensor,
-    box_crop_context: float = 0.0,
-    image_path: str = "",
-) -> torch.Tensor:
-    # box_crop_context: rate of expansion for bbox
-    # returns possibly expanded bbox xyxy as float
-
-    bbox = bbox.clone()  # do not edit bbox in place
-
-    # increase box size
-    if box_crop_context > 0.0:
-        c = box_crop_context
-        bbox = bbox.float()
-        bbox[0] -= bbox[2] * c / 2
-        bbox[1] -= bbox[3] * c / 2
-        bbox[2] += bbox[2] * c
-        bbox[3] += bbox[3] * c
-
-    if (bbox[2:] <= 1.0).any():
-        raise ValueError(
-            f"squashed image {image_path}!! The bounding box contains no pixels."
-        )
-
-    bbox[2:] = torch.clamp(bbox[2:], 2)  # set min height, width to 2 along both axes
-    bbox_xyxy = _bbox_xywh_to_xyxy(bbox, clamp_size=2)
-
-    return bbox_xyxy
-
-
-def _crop_around_box(tensor, bbox, impath: str = ""):
-    # bbox is xyxy, where the upper bound is corrected with +1
-    bbox = _clamp_box_to_image_bounds_and_round(
-        bbox,
-        image_size_hw=tensor.shape[-2:],
-    )
-    tensor = tensor[..., bbox[1] : bbox[3], bbox[0] : bbox[2]]
-    assert all(c > 0 for c in tensor.shape), f"squashed image {impath}"
-    return tensor
-
-
-def _clamp_box_to_image_bounds_and_round(
-    bbox_xyxy: torch.Tensor,
-    image_size_hw: Tuple[int, int],
-) -> torch.LongTensor:
-    bbox_xyxy = bbox_xyxy.clone()
-    bbox_xyxy[[0, 2]] = torch.clamp(bbox_xyxy[[0, 2]], 0, image_size_hw[-1])
-    bbox_xyxy[[1, 3]] = torch.clamp(bbox_xyxy[[1, 3]], 0, image_size_hw[-2])
-    if not isinstance(bbox_xyxy, torch.LongTensor):
-        bbox_xyxy = bbox_xyxy.round().long()
-    return bbox_xyxy  # pyre-ignore [7]
-
-
-def _rescale_bbox(bbox: torch.Tensor, orig_res, new_res) -> torch.Tensor:
-    assert bbox is not None
-    assert np.prod(orig_res) > 1e-8
-    # average ratio of dimensions
-    rel_size = (new_res[0] / orig_res[0] + new_res[1] / orig_res[1]) / 2.0
-    return bbox * rel_size
-
-
-def _bbox_xyxy_to_xywh(xyxy: torch.Tensor) -> torch.Tensor:
-    wh = xyxy[2:] - xyxy[:2]
-    xywh = torch.cat([xyxy[:2], wh])
-    return xywh
-
-
-def _bbox_xywh_to_xyxy(
-    xywh: torch.Tensor, clamp_size: Optional[int] = None
-) -> torch.Tensor:
-    xyxy = xywh.clone()
-    if clamp_size is not None:
-        xyxy[2:] = torch.clamp(xyxy[2:], clamp_size)
-    xyxy[2:] += xyxy[:2]
-    return xyxy
-
-
-def _safe_as_tensor(data, dtype):
-    if data is None:
-        return None
-    return torch.tensor(data, dtype=dtype)
-
-
-# NOTE this cache is per-worker; they are implemented as processes.
-# each batch is loaded and collated by a single worker;
-# since sequences tend to co-occur within batches, this is useful.
-@functools.lru_cache(maxsize=256)
-def _load_pointcloud(pcl_path: Union[str, Path], max_points: int = 0) -> Pointclouds:
-    pcl = IO().load_pointcloud(pcl_path)
-    if max_points > 0:
-        pcl = pcl.subsample(max_points)
-
-    return pcl

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