mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 10:34:29 +08:00
Switching to Torch 2.0 by default. (#1922)
* Update requirements for torch 2.0 * Update docs * Fix missing _LRScheduler import * Fix pylint issues * Update docs, drop 1.13 * Fix pip postfix version * fix docs * Allow python 1.13 (for now) * Minor fix for torch.compile * Switch to torch.compile * ruff * Fix test --------- Co-authored-by: Matt Tancik <mattt4001@gmail.com> Co-authored-by: Matthew Tancik <tancik@berkeley.edu>
This commit is contained in:
Jonáš Kulhánek
GitHub
parent
be291970db
commit
e898f562e6
@ -114,8 +114,8 @@ SHELL ["/bin/bash", "-c"]
|
||||
|
||||
# Upgrade pip and install packages.
|
||||
RUN python3.10 -m pip install --upgrade pip setuptools pathtools promise pybind11
|
||||
# Install pytorch and submodules (Currently, we still use cu116 which is the latest version for torch 1.12.1 and is compatible with CUDA 11.8).
|
||||
RUN python3.10 -m pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
|
||||
# Install pytorch and submodules
|
||||
RUN python3.10 -m pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
|
||||
# Install tynyCUDNN (we need to set the target architectures as environment variable first).
|
||||
ENV TCNN_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}
|
||||
RUN python3.10 -m pip install git+https://github.com/NVlabs/tiny-cuda-nn.git@v1.6#subdirectory=bindings/torch
|
||||
|
||||
20
README.md
20
README.md
@ -83,7 +83,7 @@ For more complex changes (e.g., running with your own data/setting up a new NeRF
|
||||
|
||||
### Prerequisites
|
||||
|
||||
You must have an NVIDIA video card with CUDA installed on the system. This library has been tested with version 11.3 of CUDA. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html)
|
||||
You must have an NVIDIA video card with CUDA installed on the system. This library has been tested with version 11.8 of CUDA. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html)
|
||||
|
||||
### Create environment
|
||||
|
||||
@ -97,19 +97,19 @@ python -m pip install --upgrade pip
|
||||
|
||||
### Dependencies
|
||||
|
||||
Install pytorch with CUDA (this repo has been tested with CUDA 11.3 and CUDA 11.7) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
|
||||
|
||||
For CUDA 11.3:
|
||||
|
||||
```bash
|
||||
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
|
||||
pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
|
||||
```
|
||||
Install pytorch with CUDA (this repo has been tested with CUDA 11.7 and CUDA 11.8) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
|
||||
|
||||
For CUDA 11.7:
|
||||
|
||||
```bash
|
||||
pip install torch==1.13.1 torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
|
||||
pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
|
||||
pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
|
||||
```
|
||||
|
||||
For CUDA 11.8:
|
||||
|
||||
```bash
|
||||
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
|
||||
pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
|
||||
```
|
||||
|
||||
|
||||
@ -4,20 +4,13 @@
|
||||
::::::{tab-set}
|
||||
:::::{tab-item} Linux
|
||||
|
||||
Install CUDA. This library has been tested with version 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
|
||||
Install CUDA. This library has been tested with version 11.8. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
|
||||
|
||||
:::::
|
||||
:::::{tab-item} Windows
|
||||
Install [Git](https://git-scm.com/downloads).
|
||||
|
||||
Install CUDA. This library and the Windows installation process has been tested with version 11.3.
|
||||
|
||||
::::{tab-set}
|
||||
:::{tab-item} 11.3 with Visual Studio 2019 v16.9
|
||||
|
||||
Install Visual Studio 2019 ver 16.9. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition). You can find older versions of Visual Studio 2019 [here](https://learn.microsoft.com/en-us/visualstudio/releases/2019/history#release-dates-and-build-numbers). The latest version, 16.11, may cause errors when installing tiny-cuda-nn.
|
||||
|
||||
Install CUDA 11.3. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html).
|
||||
Install CUDA. This library and the Windows installation process has been tested with version 11.7.
|
||||
|
||||
:::
|
||||
:::{tab-item} 11.7 with Visual Studio 2022
|
||||
@ -48,30 +41,39 @@ python -m pip install --upgrade pip
|
||||
### pytorch
|
||||
|
||||
::::{tab-set}
|
||||
:::{tab-item} Torch 1.12.1 with CUDA 11.3
|
||||
:::{tab-item} Torch 2.0.1 with CUDA 11.8
|
||||
|
||||
- To install 1.12.1 with CUDA 11.3:
|
||||
- To install 2.0.1 with CUDA 11.8:
|
||||
|
||||
```bash
|
||||
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
|
||||
```
|
||||
|
||||
:::
|
||||
:::{tab-item} Torch 1.13.1 with CUDA 11.7
|
||||
|
||||
- To install 1.13.1 with CUDA 11.7:
|
||||
|
||||
Note that if a pytorch version prior to 1.13 is installed,
|
||||
Note that if a pytorch version prior to 2.0.1 is installed,
|
||||
it should be uninstalled first to avoid upgrade issues (e.g. with functorch)
|
||||
|
||||
```bash
|
||||
pip uninstall torch torchvision functorch
|
||||
```
|
||||
|
||||
Install pytorch 1.13.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
|
||||
Install pytorch 2.0.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
|
||||
|
||||
```bash
|
||||
pip install torch==1.13.1 torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117
|
||||
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
|
||||
```
|
||||
|
||||
:::
|
||||
:::{tab-item} Torch 2.0.1 with CUDA 11.7
|
||||
|
||||
- To install 2.0.1 with CUDA 11.7:
|
||||
|
||||
Note that if a pytorch version prior to 2.0.1 is installed,
|
||||
it should be uninstalled first to avoid upgrade issues (e.g. with functorch)
|
||||
|
||||
```bash
|
||||
pip uninstall torch torchvision functorch
|
||||
```
|
||||
|
||||
Install pytorch 2.0.1 with CUDA and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn)
|
||||
|
||||
```bash
|
||||
pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
|
||||
```
|
||||
|
||||
:::
|
||||
@ -194,7 +196,6 @@ K80 | 37
|
||||
|
||||
- [TinyCUDA installation errors out with cuda mismatch](tiny-cuda-mismatch-error)
|
||||
- [TinyCUDA installation errors out with no CUDA toolset found](tiny-cuda-integration-error)
|
||||
- [TinyCUDA installation errors out with syntax errors](tiny-cuda-syntax-error)
|
||||
- [Installation errors, File "setup.py" not found](pip-install-error)
|
||||
- [Runtime errors, "len(sources) > 0".](cuda-sources-error)
|
||||
|
||||
@ -221,7 +222,7 @@ While installing tiny-cuda on Windows, you run into: `No CUDA toolset found.`
|
||||
|
||||
**Solution**:
|
||||
|
||||
Confirm that you have Visual Studio installed (CUDA 11.3 is only compatible up to 2019 ver 16.9, not 16.11 or 2022).
|
||||
Confirm that you have Visual Studio installed.
|
||||
|
||||
Make sure CUDA Visual Studio integration is enabled. This should be done automatically by the CUDA installer if it is run after Visual Studio is installed. You can also manually enable integration.
|
||||
|
||||
@ -231,7 +232,7 @@ Make sure CUDA Visual Studio integration is enabled. This should be done automat
|
||||
To manually enable integration for Visual Studio 2019, copy all 4 files from
|
||||
|
||||
```
|
||||
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\extras\visual_studio_integration\MSBuildExtensions
|
||||
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\extras\visual_studio_integration\MSBuildExtensions
|
||||
```
|
||||
|
||||
to
|
||||
@ -246,7 +247,7 @@ C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\MSBuild\Microsoft
|
||||
To manually enable integration for Visual Studio 2022, copy all 4 files from
|
||||
|
||||
```
|
||||
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\extras\visual_studio_integration\MSBuildExtensions
|
||||
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\extras\visual_studio_integration\MSBuildExtensions
|
||||
```
|
||||
|
||||
to
|
||||
@ -260,20 +261,6 @@ C:\Program Files\Microsoft Visual Studio\2022\[Community, Professional, Enterpri
|
||||
|
||||
<br />
|
||||
|
||||
(tiny-cuda-syntax-error)=
|
||||
|
||||
**(Windows) TinyCUDA installation errors out with syntax errors**
|
||||
|
||||
While installing tiny-cuda on Windows, you run into: `Expected a "("`
|
||||
|
||||
**Solution**:
|
||||
|
||||
If you are using CUDA 11.3 with Visual Studio 2019, confirm that Visual Studio is on version 16.9 and not 16.11.
|
||||
|
||||
If you are using a different version of CUDA and Visual Studio, try switching to CUDA 11.3 with Visual Studio 2019 ver 16.9.
|
||||
|
||||
<br />
|
||||
|
||||
(pip-install-error)=
|
||||
|
||||
**Installation errors, File "setup.py" not found**
|
||||
|
||||
@ -25,6 +25,8 @@ from jaxtyping import Float
|
||||
from numpy.typing import NDArray
|
||||
from torch import Tensor
|
||||
|
||||
from nerfstudio.utils.misc import torch_compile
|
||||
|
||||
_EPS = np.finfo(float).eps * 4.0
|
||||
|
||||
|
||||
@ -340,7 +342,6 @@ def get_distortion_params(
|
||||
return torch.Tensor([k1, k2, k3, k4, p1, p2])
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def _compute_residual_and_jacobian(
|
||||
x: torch.Tensor,
|
||||
y: torch.Tensor,
|
||||
@ -406,7 +407,7 @@ def _compute_residual_and_jacobian(
|
||||
return fx, fy, fx_x, fx_y, fy_x, fy_y
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
@torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
|
||||
def radial_and_tangential_undistort(
|
||||
coords: torch.Tensor,
|
||||
distortion_params: torch.Tensor,
|
||||
|
||||
@ -24,6 +24,14 @@ from torch.optim import Optimizer, lr_scheduler
|
||||
|
||||
from nerfstudio.configs.base_config import InstantiateConfig
|
||||
|
||||
try:
|
||||
from torch.optim.lr_scheduler import ( # pylint: disable=ungrouped-imports
|
||||
LRScheduler,
|
||||
)
|
||||
except ImportError:
|
||||
# Backward compatibility for PyTorch 1.x
|
||||
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
|
||||
|
||||
|
||||
@dataclass
|
||||
class SchedulerConfig(InstantiateConfig):
|
||||
@ -43,7 +51,7 @@ class Scheduler:
|
||||
self.config = config
|
||||
|
||||
@abstractmethod
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
|
||||
"""Abstract method that returns a scheduler object.
|
||||
|
||||
Args:
|
||||
@ -73,7 +81,7 @@ class MultiStepScheduler(Scheduler):
|
||||
|
||||
config: MultiStepSchedulerConfig
|
||||
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
|
||||
scheduler = lr_scheduler.MultiStepLR(
|
||||
optimizer=optimizer,
|
||||
milestones=self.config.milestones,
|
||||
@ -107,7 +115,7 @@ class ExponentialDecayScheduler(Scheduler):
|
||||
|
||||
config: ExponentialDecaySchedulerConfig
|
||||
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
|
||||
if self.config.lr_final is None:
|
||||
lr_final = lr_init
|
||||
else:
|
||||
@ -154,7 +162,7 @@ class CosineDecayScheduler(Scheduler):
|
||||
|
||||
config: CosineDecaySchedulerConfig
|
||||
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
|
||||
def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> LRScheduler:
|
||||
def func(step):
|
||||
if step < self.config.warm_up_end:
|
||||
learning_factor = step / self.config.warm_up_end
|
||||
|
||||
@ -22,6 +22,7 @@ import torch
|
||||
from jaxtyping import Float
|
||||
from torch import Tensor, nn
|
||||
|
||||
from nerfstudio.utils.misc import torch_compile
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
from nerfstudio.data.scene_box import SceneBox
|
||||
|
||||
@ -108,7 +109,7 @@ class AABBBoxCollider(SceneCollider):
|
||||
return ray_bundle
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
@torch_compile(dynamic=True, mode="reduce-overhead")
|
||||
def _intersect_with_sphere(
|
||||
rays_o: torch.Tensor, rays_d: torch.Tensor, center: torch.Tensor, radius: float = 1.0, near_plane: float = 0.0
|
||||
):
|
||||
|
||||
@ -21,6 +21,8 @@ import torch
|
||||
from jaxtyping import Bool, Float
|
||||
from torch import Tensor
|
||||
|
||||
from nerfstudio.utils.misc import torch_compile
|
||||
|
||||
|
||||
def components_from_spherical_harmonics(
|
||||
levels: int, directions: Float[Tensor, "*batch 3"]
|
||||
@ -197,7 +199,7 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
|
||||
return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
@torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
|
||||
def intersect_aabb(
|
||||
origins: torch.Tensor,
|
||||
directions: torch.Tensor,
|
||||
@ -222,11 +224,11 @@ def intersect_aabb(
|
||||
tx_min = (aabb[:3] - origins) / directions
|
||||
tx_max = (aabb[3:] - origins) / directions
|
||||
|
||||
t_min = torch.min(tx_min, tx_max)
|
||||
t_max = torch.max(tx_min, tx_max)
|
||||
t_min = torch.stack((tx_min, tx_max)).amin(dim=0)
|
||||
t_max = torch.stack((tx_min, tx_max)).amax(dim=0)
|
||||
|
||||
t_min = torch.max(t_min, dim=-1).values
|
||||
t_max = torch.min(t_max, dim=-1).values
|
||||
t_min = t_min.amax(dim=-1)
|
||||
t_max = t_max.amin(dim=-1)
|
||||
|
||||
t_min = torch.clamp(t_min, min=0, max=max_bound)
|
||||
t_max = torch.clamp(t_max, min=0, max=max_bound)
|
||||
@ -322,3 +324,4 @@ def normalized_depth_scale_and_shift(
|
||||
shift[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
|
||||
|
||||
return scale, shift
|
||||
return scale, shift
|
||||
|
||||
@ -155,3 +155,13 @@ def strtobool(val) -> bool:
|
||||
FMI https://stackoverflow.com/a/715468
|
||||
"""
|
||||
return val.lower() in ("yes", "y", "true", "t", "on", "1")
|
||||
|
||||
|
||||
if hasattr(torch, "compile"):
|
||||
# PyTorch 2.0+
|
||||
torch_compile = torch.compile
|
||||
else:
|
||||
# PyTorch 1.x
|
||||
def torch_compile(*args, **kwargs):
|
||||
"""Fall back to jit script if using pytorch 1.x, args ignored"""
|
||||
return torch.jit.script
|
||||
|
||||
@ -70,11 +70,10 @@ class TensorDataclass:
|
||||
This will 1) find the broadcasted shape and 2) broadcast all fields to this shape 3)
|
||||
set _shape to be the broadcasted shape.
|
||||
"""
|
||||
if self._field_custom_dimensions is not None:
|
||||
for k, v in self._field_custom_dimensions.items():
|
||||
assert (
|
||||
isinstance(v, int) and v > 1
|
||||
), f"Custom dimensions must be an integer greater than 1, since 1 is the default, received {k}: {v}"
|
||||
for k, v in self._field_custom_dimensions.items():
|
||||
assert (
|
||||
isinstance(v, int) and v > 1
|
||||
), f"Custom dimensions must be an integer greater than 1, since 1 is the default, received {k}: {v}"
|
||||
|
||||
# Shim to prevent pyright from narrowing `self` to DataclassInstance.
|
||||
self_dc = self
|
||||
@ -130,6 +129,7 @@ class TensorDataclass:
|
||||
if isinstance(v, torch.Tensor):
|
||||
# If custom dimension key, then we need to
|
||||
if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions:
|
||||
# pylint: disable=unsubscriptable-object
|
||||
new_dict[k] = v.broadcast_to(
|
||||
(
|
||||
*batch_shape,
|
||||
@ -326,7 +326,6 @@ class TensorDataclass:
|
||||
# This is the case when we have a custom dimensions tensor
|
||||
elif (
|
||||
isinstance(v, torch.Tensor)
|
||||
and isinstance(self._field_custom_dimensions, dict)
|
||||
and f in self._field_custom_dimensions
|
||||
and custom_tensor_dims_fn is not None
|
||||
):
|
||||
|
||||
@ -20,7 +20,6 @@ dependencies = [
|
||||
"tyro>=0.3.31",
|
||||
"gdown>=4.6.0",
|
||||
"ninja>=1.10",
|
||||
"functorch>=0.2.1",
|
||||
"h5py>=2.9.0",
|
||||
"imageio>=2.21.1",
|
||||
'importlib-metadata>=6.0.0; python_version < "3.10"',
|
||||
@ -48,9 +47,9 @@ dependencies = [
|
||||
"rich>=12.5.1",
|
||||
"scikit-image>=0.19.3",
|
||||
"tensorboard>=2.13.0",
|
||||
"torch>=1.12.1,<2.0.0",
|
||||
"torch>=1.13.1",
|
||||
"torchvision>=0.14.1",
|
||||
"torchmetrics[image]>=0.9.3",
|
||||
"torchvision>=0.13.0",
|
||||
"typing_extensions>=4.4.0",
|
||||
"viser>=0.0.5",
|
||||
"nuscenes-devkit>=1.1.1",
|
||||
@ -115,7 +114,7 @@ ns-dev-sync-viser-message-defs = "nerfstudio.scripts.viewer.sync_viser_message_d
|
||||
[options]
|
||||
# equivalent to using --extra-index-url with pip, which is needed for specifying the CUDA version torch and torchvision
|
||||
dependency_links = [
|
||||
"https://download.pytorch.org/whl/cu113"
|
||||
"https://download.pytorch.org/whl/cu118"
|
||||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user