diff --git a/docs/renderer.html b/docs/renderer.html index bf48cd3f..c6f8d0f0 100644 --- a/docs/renderer.html +++ b/docs/renderer.html @@ -75,7 +75,7 @@

In order to experiment with different approaches, we wanted a modular implementation that is easy to use and extend, and supports heterogeneous batching. Taking inspiration from existing work [1, 2], we have created a new, modular, differentiable renderer with parallel implementations in PyTorch, C++ and CUDA, as well as comprehensive documentation and tests, with the aim of helping to further research in this field.

Our implementation decouples the rasterization and shading steps of rendering. The core rasterization step (based on [2]) returns several intermediate variables and has an optimized implementation in CUDA. The rest of the pipeline is implemented purely in PyTorch, and is designed to be customized and extended. With this approach, the PyTorch3D differentiable renderer can be imported as a library.

Get started

-

To learn about more the implementation and start using the renderer refer to getting started with renderer, which also contains the architecture overview and coordinate transformation conventions.

+

To learn about more the implementation and start using the renderer refer to getting started with renderer, which also contains the architecture overview and coordinate transformation conventions.

Tech Report

For an in depth explanation of the renderer design, key features and benchmarks please refer to the PyTorch3D Technical Report on ArXiv: Accelerating 3D Deep Learning with PyTorch3D, for the pulsar backend see here: Fast Differentiable Raycasting for Neural Rendering using Sphere-based Representations.

@@ -109,4 +109,4 @@ total_memory = memory_forward_pass + memory_backward_pass

[6] Yifan et al, 'Differentiable Surface Splatting for Point-based Geometry Processing', SIGGRAPH Asia 2019

[7] Loubet et al, 'Reparameterizing Discontinuous Integrands for Differentiable Rendering', SIGGRAPH Asia 2019

[8] Chen et al, 'Learning to Predict 3D Objects with an Interpolation-based Differentiable Renderer', NeurIPS 2019

-
Last updated by Christoph Lassner
← Plotly VisualizationGetting Started →
\ No newline at end of file +
Last updated by Krzysztof Chalupka
← Plotly VisualizationGetting Started →
\ No newline at end of file diff --git a/docs/renderer/index.html b/docs/renderer/index.html index bf48cd3f..c6f8d0f0 100644 --- a/docs/renderer/index.html +++ b/docs/renderer/index.html @@ -75,7 +75,7 @@

In order to experiment with different approaches, we wanted a modular implementation that is easy to use and extend, and supports heterogeneous batching. Taking inspiration from existing work [1, 2], we have created a new, modular, differentiable renderer with parallel implementations in PyTorch, C++ and CUDA, as well as comprehensive documentation and tests, with the aim of helping to further research in this field.

Our implementation decouples the rasterization and shading steps of rendering. The core rasterization step (based on [2]) returns several intermediate variables and has an optimized implementation in CUDA. The rest of the pipeline is implemented purely in PyTorch, and is designed to be customized and extended. With this approach, the PyTorch3D differentiable renderer can be imported as a library.

Get started

-

To learn about more the implementation and start using the renderer refer to getting started with renderer, which also contains the architecture overview and coordinate transformation conventions.

+

To learn about more the implementation and start using the renderer refer to getting started with renderer, which also contains the architecture overview and coordinate transformation conventions.

Tech Report

For an in depth explanation of the renderer design, key features and benchmarks please refer to the PyTorch3D Technical Report on ArXiv: Accelerating 3D Deep Learning with PyTorch3D, for the pulsar backend see here: Fast Differentiable Raycasting for Neural Rendering using Sphere-based Representations.

@@ -109,4 +109,4 @@ total_memory = memory_forward_pass + memory_backward_pass

[6] Yifan et al, 'Differentiable Surface Splatting for Point-based Geometry Processing', SIGGRAPH Asia 2019

[7] Loubet et al, 'Reparameterizing Discontinuous Integrands for Differentiable Rendering', SIGGRAPH Asia 2019

[8] Chen et al, 'Learning to Predict 3D Objects with an Interpolation-based Differentiable Renderer', NeurIPS 2019

-
Last updated by Christoph Lassner
← Plotly VisualizationGetting Started →
\ No newline at end of file +
Last updated by Krzysztof Chalupka
← Plotly VisualizationGetting Started →
\ No newline at end of file diff --git a/files/bundle_adjustment.ipynb b/files/bundle_adjustment.ipynb index 59bad817..91640c9b 100644 --- a/files/bundle_adjustment.ipynb +++ b/files/bundle_adjustment.ipynb @@ -89,7 +89,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/bundle_adjustment.py b/files/bundle_adjustment.py index 7f197e63..7bf593af 100644 --- a/files/bundle_adjustment.py +++ b/files/bundle_adjustment.py @@ -49,7 +49,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/camera_position_optimization_with_differentiable_rendering.ipynb b/files/camera_position_optimization_with_differentiable_rendering.ipynb index f3b66485..e8e81e0a 100644 --- a/files/camera_position_optimization_with_differentiable_rendering.ipynb +++ b/files/camera_position_optimization_with_differentiable_rendering.ipynb @@ -76,7 +76,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/camera_position_optimization_with_differentiable_rendering.py b/files/camera_position_optimization_with_differentiable_rendering.py index d6b0df46..92c656c1 100644 --- a/files/camera_position_optimization_with_differentiable_rendering.py +++ b/files/camera_position_optimization_with_differentiable_rendering.py @@ -36,7 +36,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/dataloaders_ShapeNetCore_R2N2.ipynb b/files/dataloaders_ShapeNetCore_R2N2.ipynb index 1110bdc4..728a103a 100644 --- a/files/dataloaders_ShapeNetCore_R2N2.ipynb +++ b/files/dataloaders_ShapeNetCore_R2N2.ipynb @@ -51,7 +51,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/dataloaders_ShapeNetCore_R2N2.py b/files/dataloaders_ShapeNetCore_R2N2.py index b7429dc8..928115a9 100644 --- a/files/dataloaders_ShapeNetCore_R2N2.py +++ b/files/dataloaders_ShapeNetCore_R2N2.py @@ -31,7 +31,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/deform_source_mesh_to_target_mesh.ipynb b/files/deform_source_mesh_to_target_mesh.ipynb index 172b0046..f6f2d7d3 100644 --- a/files/deform_source_mesh_to_target_mesh.ipynb +++ b/files/deform_source_mesh_to_target_mesh.ipynb @@ -90,7 +90,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/deform_source_mesh_to_target_mesh.py b/files/deform_source_mesh_to_target_mesh.py index bafccf16..5c18f91c 100644 --- a/files/deform_source_mesh_to_target_mesh.py +++ b/files/deform_source_mesh_to_target_mesh.py @@ -48,7 +48,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/fit_simple_neural_radiance_field.ipynb b/files/fit_simple_neural_radiance_field.ipynb index 6490c527..56b7876e 100644 --- a/files/fit_simple_neural_radiance_field.ipynb +++ b/files/fit_simple_neural_radiance_field.ipynb @@ -56,7 +56,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", @@ -815,7 +815,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## 5. Visualizing the optimized neural radiance field\n", + "## 6. Visualizing the optimized neural radiance field\n", "\n", "Finally, we visualize the neural radiance field by rendering from multiple viewpoints that rotate around the volume's y-axis." ] @@ -865,7 +865,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## 6. Conclusion\n", + "## 7. Conclusion\n", "\n", "In this tutorial, we have shown how to optimize an implicit representation of a scene such that the renders of the scene from known viewpoints match the observed images for each viewpoint. The rendering was carried out using the PyTorch3D's implicit function renderer composed of either a `MonteCarloRaysampler` or `NDCMultinomialRaysampler`, and an `EmissionAbsorptionRaymarcher`." ] diff --git a/files/fit_simple_neural_radiance_field.py b/files/fit_simple_neural_radiance_field.py index a4f02561..dee73a94 100644 --- a/files/fit_simple_neural_radiance_field.py +++ b/files/fit_simple_neural_radiance_field.py @@ -42,7 +42,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ @@ -744,7 +744,7 @@ for iteration in range(n_iter): ) -# ## 5. Visualizing the optimized neural radiance field +# ## 6. Visualizing the optimized neural radiance field # # Finally, we visualize the neural radiance field by rendering from multiple viewpoints that rotate around the volume's y-axis. @@ -786,6 +786,6 @@ image_grid(rotating_nerf_frames.clamp(0., 1.).cpu().numpy(), rows=3, cols=5, rgb plt.show() -# ## 6. Conclusion +# ## 7. Conclusion # # In this tutorial, we have shown how to optimize an implicit representation of a scene such that the renders of the scene from known viewpoints match the observed images for each viewpoint. The rendering was carried out using the PyTorch3D's implicit function renderer composed of either a `MonteCarloRaysampler` or `NDCMultinomialRaysampler`, and an `EmissionAbsorptionRaymarcher`. diff --git a/files/fit_textured_mesh.ipynb b/files/fit_textured_mesh.ipynb index e3a3bca0..954884b0 100644 --- a/files/fit_textured_mesh.ipynb +++ b/files/fit_textured_mesh.ipynb @@ -68,7 +68,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", @@ -118,7 +118,7 @@ "from pytorch3d.structures import Meshes\n", "from pytorch3d.renderer import (\n", " look_at_view_transform,\n", - " OpenGLPerspectiveCameras, \n", + " FoVPerspectiveCameras, \n", " PointLights, \n", " DirectionalLights, \n", " Materials, \n", @@ -304,11 +304,11 @@ "# broadcasting. So we can view the camera from the a distance of dist=2.7, and \n", "# then specify elevation and azimuth angles for each viewpoint as tensors. \n", "R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)\n", - "cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)\n", + "cameras = FoVPerspectiveCameras(device=device, R=R, T=T)\n", "\n", "# We arbitrarily choose one particular view that will be used to visualize \n", "# results\n", - "camera = OpenGLPerspectiveCameras(device=device, R=R[None, 1, ...], \n", + "camera = FoVPerspectiveCameras(device=device, R=R[None, 1, ...], \n", " T=T[None, 1, ...]) \n", "\n", "# Define the settings for rasterization and shading. Here we set the output \n", @@ -351,7 +351,7 @@ "# Our multi-view cow dataset will be represented by these 2 lists of tensors,\n", "# each of length num_views.\n", "target_rgb = [target_images[i, ..., :3] for i in range(num_views)]\n", - "target_cameras = [OpenGLPerspectiveCameras(device=device, R=R[None, i, ...], \n", + "target_cameras = [FoVPerspectiveCameras(device=device, R=R[None, i, ...], \n", " T=T[None, i, ...]) for i in range(num_views)]" ] }, @@ -708,6 +708,7 @@ " image_size=128, \n", " blur_radius=np.log(1. / 1e-4 - 1.)*sigma, \n", " faces_per_pixel=50, \n", + " perspective_correct=False, \n", ")\n", "\n", "# Differentiable soft renderer using per vertex RGB colors for texture\n", diff --git a/files/fit_textured_mesh.py b/files/fit_textured_mesh.py index b949e4e1..b10be54d 100644 --- a/files/fit_textured_mesh.py +++ b/files/fit_textured_mesh.py @@ -31,7 +31,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ @@ -74,7 +74,7 @@ from pytorch3d.loss import ( from pytorch3d.structures import Meshes from pytorch3d.renderer import ( look_at_view_transform, - OpenGLPerspectiveCameras, + FoVPerspectiveCameras, PointLights, DirectionalLights, Materials, @@ -185,11 +185,11 @@ lights = PointLights(device=device, location=[[0.0, 0.0, -3.0]]) # broadcasting. So we can view the camera from the a distance of dist=2.7, and # then specify elevation and azimuth angles for each viewpoint as tensors. R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim) -cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T) +cameras = FoVPerspectiveCameras(device=device, R=R, T=T) # We arbitrarily choose one particular view that will be used to visualize # results -camera = OpenGLPerspectiveCameras(device=device, R=R[None, 1, ...], +camera = FoVPerspectiveCameras(device=device, R=R[None, 1, ...], T=T[None, 1, ...]) # Define the settings for rasterization and shading. Here we set the output @@ -232,7 +232,7 @@ target_images = renderer(meshes, cameras=cameras, lights=lights) # Our multi-view cow dataset will be represented by these 2 lists of tensors, # each of length num_views. target_rgb = [target_images[i, ..., :3] for i in range(num_views)] -target_cameras = [OpenGLPerspectiveCameras(device=device, R=R[None, i, ...], +target_cameras = [FoVPerspectiveCameras(device=device, R=R[None, i, ...], T=T[None, i, ...]) for i in range(num_views)] @@ -461,6 +461,7 @@ raster_settings_soft = RasterizationSettings( image_size=128, blur_radius=np.log(1. / 1e-4 - 1.)*sigma, faces_per_pixel=50, + perspective_correct=False, ) # Differentiable soft renderer using per vertex RGB colors for texture diff --git a/files/fit_textured_volume.ipynb b/files/fit_textured_volume.ipynb index bae176bf..9c3edf42 100644 --- a/files/fit_textured_volume.ipynb +++ b/files/fit_textured_volume.ipynb @@ -47,7 +47,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/fit_textured_volume.py b/files/fit_textured_volume.py index fbf5c8f1..7cef3430 100644 --- a/files/fit_textured_volume.py +++ b/files/fit_textured_volume.py @@ -32,7 +32,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/implicitron_config_system.ipynb b/files/implicitron_config_system.ipynb new file mode 100644 index 00000000..365acf19 --- /dev/null +++ b/files/implicitron_config_system.ipynb @@ -0,0 +1,1252 @@ +{ + "metadata": { + "dataExplorerConfig": {}, + "bento_stylesheets": { + "bento/extensions/flow/main.css": true, + "bento/extensions/kernel_selector/main.css": true, + "bento/extensions/kernel_ui/main.css": true, + "bento/extensions/new_kernel/main.css": true, + "bento/extensions/system_usage/main.css": true, + "bento/extensions/theme/main.css": true + }, + "kernelspec": { + "display_name": "pytorch3d", + "language": "python", + "name": "bento_kernel_pytorch3d", + "metadata": { + "kernel_name": "bento_kernel_pytorch3d", + "nightly_builds": true, + "fbpkg_supported": true, + "cinder_runtime": false, + "is_prebuilt": true + } + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + }, + "last_server_session_id": "d6b46f14-cee7-44c1-8c51-39a38a4ea4c2", + "last_kernel_id": "90755407-3729-46f4-ab67-ff2cb1daa5cb", + "last_base_url": "https://9177.od.fbinfra.net:443/", + "last_msg_id": "f61034eb-826226915ad9548ffbe495ba_6317", + "captumWidgetMessage": {}, + "outputWidgetContext": {} + }, + "nbformat": 4, + "nbformat_minor": 2, + "cells": [ + { + "cell_type": "code", + "metadata": { + "originalKey": "f0af2d90-cb21-4ab4-b4cb-0fd00dbfb77b", + "showInput": true, + "customInput": null, + "customOutput": null + }, + "source": [ + "# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved." + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "4e15bfa2-5404-40d0-98b6-eb2732c8b72b", + "showInput": false, + "customInput": null + }, + "source": [ + "# Implicitron's config system" + ], + "attachments": {} + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "287be985-423d-42e0-a2af-1e8c585e723c", + "showInput": false, + "customInput": null + }, + "source": [ + "Implicitron's components are all based on a unified hierarchical configuration system. \n", + "This allows configurable variables and all defaults to be defined separately for each new component.\n", + "All configs relevant to an experiment are then automatically composed into a single configuration file that fully specifies the experiment.\n", + "An especially important feature is extension points where users can insert their own sub-classes of Implicitron's base components.\n", + "\n", + "The file which defines this system is [here](https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/implicitron/tools/config.py) in the PyTorch3D repo.\n", + "The Implicitron volumes tutorial contains a simple example of using the config system.\n", + "This tutorial provides detailed hands-on experience in using and modifying Implicitron's configurable components.\n", + "" + ], + "attachments": {} + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "fde300a2-99cb-4d52-9d5b-4464a2083e0b", + "showInput": false, + "customInput": null + }, + "source": [ + "## 0. Install and import modules\n", + "\n", + "Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell:" + ] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "ad6e94a7-e114-43d3-b038-a5210c7d34c9", + "showInput": true, + "customInput": null, + "customOutput": null + }, + "source": [ + "import os\n", + "import sys\n", + "import torch\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.12.\") 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", + " !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'" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "609896c0-9e2e-4716-b074-b565f0170e32", + "showInput": false, + "customInput": null + }, + "source": [ + "Ensure omegaconf is installed. If not, run this cell. (It should not be necessary to restart the runtime.)" + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "d1c1851e-b9f2-4236-93c3-19aa4d63041c", + "showInput": true, + "customInput": null, + "customOutput": null + }, + "source": [ + "!pip install omegaconf" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "collapsed": false, + "originalKey": "5ac7ef23-b74c-46b2-b8d3-799524d7ba4f", + "code_folding": [], + "hidden_ranges": [], + "requestMsgId": "5ac7ef23-b74c-46b2-b8d3-799524d7ba4f", + "customOutput": null, + "executionStartTime": 1659465468717, + "executionStopTime": 1659465468738 + }, + "source": [ + "from dataclasses import dataclass\n", + "from typing import Optional, Tuple\n", + "\n", + "import torch\n", + "from omegaconf import DictConfig, OmegaConf\n", + "from pytorch3d.implicitron.tools.config import (\n", + " Configurable,\n", + " ReplaceableBase,\n", + " expand_args_fields,\n", + " get_default_args,\n", + " registry,\n", + " run_auto_creation,\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "a638bf90-eb6b-424d-b53d-eae11954a717", + "showInput": false, + "customInput": null + }, + "source": [ + "## 1. Introducing dataclasses \n", + "\n", + "[Type hints](https://docs.python.org/3/library/typing.html) give a taxonomy of types in Python. [Dataclasses](https://docs.python.org/3/library/dataclasses.html) let you create a class based on a list of members which have names, types and possibly default values. The `__init__` function is created automatically, and calls a `__post_init__` function if present as a final step. For example" + ] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "71eaad5e-e198-492e-8610-24b0da9dd4ae", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "71eaad5e-e198-492e-8610-24b0da9dd4ae", + "customOutput": null, + "executionStartTime": 1659454972732, + "executionStopTime": 1659454972739 + }, + "source": [ + "@dataclass\n", + "class MyDataclass:\n", + " a: int\n", + " b: int = 8\n", + " c: Optional[Tuple[int, ...]] = None\n", + "\n", + " def __post_init__(self):\n", + " print(f\"created with a = {self.a}\")\n", + " self.d = 2 * self.b" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "83202a18-a3d3-44ec-a62d-b3360a302645", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "83202a18-a3d3-44ec-a62d-b3360a302645", + "customOutput": null, + "executionStartTime": 1659454973051, + "executionStopTime": 1659454973077 + }, + "source": [ + "my_dataclass_instance = MyDataclass(a=18)\n", + "assert my_dataclass_instance.d == 16" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "b67ccb9f-dc6c-4994-9b99-b5a1bcfebd70", + "showInput": false, + "customInput": null + }, + "source": [ + "👷 Note that the `dataclass` decorator here is function which modifies the definition of the class itself.\n", + "It runs immediately after the definition.\n", + "Our config system requires that implicitron library code contains classes whose modified versions need to be aware of user-defined implementations.\n", + "Therefore we need the modification of the class to be delayed. We don't use a decorator.\n", + "" + ], + "attachments": {} + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "3e90f664-99df-4387-9c45-a1ad7939ef3a", + "showInput": false, + "customInput": null + }, + "source": [ + "## 2. Introducing omegaconf and OmegaConf.structured\n", + "\n", + "The [omegaconf](https://github.com/omry/omegaconf/) library provides a DictConfig class which is like a `dict` with str keys, but with extra features for ease-of-use as a configuration system." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "81c73c9b-27ee-4aab-b55e-fb0dd67fe174", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "81c73c9b-27ee-4aab-b55e-fb0dd67fe174", + "customOutput": null, + "executionStartTime": 1659451341683, + "executionStopTime": 1659451341690 + }, + "source": [ + "dc = DictConfig({\"a\": 2, \"b\": True, \"c\": None, \"d\": \"hello\"})\n", + "assert dc.a == dc[\"a\"] == 2" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "3b5b76a9-4b76-4784-96ff-2a1212e48e48", + "showInput": false, + "customInput": null + }, + "source": [ + "OmegaConf has a serialization to and from yaml. The [Hydra](https://hydra.cc/) library relies on this for its configuration files." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "d7a25ec1-caea-46bc-a1da-4b1f040c4b61", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "d7a25ec1-caea-46bc-a1da-4b1f040c4b61", + "customOutput": null, + "executionStartTime": 1659451411835, + "executionStopTime": 1659451411936 + }, + "source": [ + "print(OmegaConf.to_yaml(dc))\n", + "assert OmegaConf.create(OmegaConf.to_yaml(dc)) == dc" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "777fecdd-8bf6-4fd8-827b-cb8af5477fa8", + "showInput": false, + "customInput": null + }, + "source": [ + "OmegaConf.structured provides a DictConfig from a dataclass or instance of a dataclass. Unlike a normal DictConfig, it is type-checked and only known keys can be added." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "de36efb4-0b08-4fb8-bb3a-be1b2c0cd162", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "de36efb4-0b08-4fb8-bb3a-be1b2c0cd162", + "customOutput": null, + "executionStartTime": 1659455098879, + "executionStopTime": 1659455098900 + }, + "source": [ + "structured = OmegaConf.structured(MyDataclass)\n", + "assert isinstance(structured, DictConfig)\n", + "print(structured)\n", + "print()\n", + "print(OmegaConf.to_yaml(structured))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "be4446da-e536-4139-9ba3-37669a5b5e61", + "showInput": false, + "customInput": null + }, + "source": [ + "`structured` knows it is missing a value for `a`." + ], + "attachments": {} + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "864811e8-1a75-4932-a85e-f681b0541ae9", + "showInput": false, + "customInput": null + }, + "source": [ + "Such an object has members compatible with the dataclass, so an initialisation can be performed as follows." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "eb88aaa0-c22f-4ffb-813a-ca957b490acb", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "eb88aaa0-c22f-4ffb-813a-ca957b490acb", + "customOutput": null, + "executionStartTime": 1659455580491, + "executionStopTime": 1659455580501 + }, + "source": [ + "structured.a = 21\n", + "my_dataclass_instance2 = MyDataclass(**structured)\n", + "print(my_dataclass_instance2)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "2d08c81c-9d18-4de9-8464-0da2d89f94f3", + "showInput": false, + "customInput": null + }, + "source": [ + "You can also call OmegaConf.structured on an instance." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "5e469bac-32a4-475d-9c09-8b64ba3f2155", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "5e469bac-32a4-475d-9c09-8b64ba3f2155", + "customOutput": null, + "executionStartTime": 1659455594700, + "executionStopTime": 1659455594737 + }, + "source": [ + "structured_from_instance = OmegaConf.structured(my_dataclass_instance)\n", + "my_dataclass_instance3 = MyDataclass(**structured_from_instance)\n", + "print(my_dataclass_instance3)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "2ed559e3-8552-465a-938f-30c72a321184", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "2ed559e3-8552-465a-938f-30c72a321184", + "customOutput": null, + "executionStartTime": 1659452594203, + "executionStopTime": 1659452594333 + }, + "source": [ + "## 3. Our approach to OmegaConf.structured\n", + "\n", + "We provide functions which are equivalent to `OmegaConf.structured` but support more features. \n", + "To achieve the above using our functions, the following is used.\n", + "Note that we indicate configurable classes using a special base class `Configurable`, not a decorator." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "9888afbd-e617-4596-ab7a-fc1073f58656", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "9888afbd-e617-4596-ab7a-fc1073f58656", + "customOutput": null, + "executionStartTime": 1659454053323, + "executionStopTime": 1659454061629 + }, + "source": [ + "class MyConfigurable(Configurable):\n", + " a: int\n", + " b: int = 8\n", + " c: Optional[Tuple[int, ...]] = None\n", + "\n", + " def __post_init__(self):\n", + " print(f\"created with a = {self.a}\")\n", + " self.d = 2 * self.b" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "e43155b4-3da5-4df1-a2f5-da1d0369eec9", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "e43155b4-3da5-4df1-a2f5-da1d0369eec9", + "customOutput": null, + "executionStartTime": 1659454784912, + "executionStopTime": 1659454784928 + }, + "source": [ + "# expand_args_fields must be called on an object before it is instantiated.\n", + "# A warning is raised if this is missed, but it is possible to not notice the warning.\n", + "# It modifies the class like @dataclass\n", + "expand_args_fields(MyConfigurable)\n", + "my_configurable_instance = MyConfigurable(a=18)\n", + "assert my_configurable_instance.d == 16" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "96eaae18-dce4-4ee1-b451-1466fea51b9f", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "96eaae18-dce4-4ee1-b451-1466fea51b9f", + "customOutput": null, + "executionStartTime": 1659460669541, + "executionStopTime": 1659460669566 + }, + "source": [ + "# get_default_args calls expand_args_fields automatically\n", + "our_structured = get_default_args(MyConfigurable)\n", + "assert isinstance(our_structured, DictConfig)\n", + "print(OmegaConf.to_yaml(our_structured))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "359f7925-68de-42cd-bd34-79a099b1c210", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "359f7925-68de-42cd-bd34-79a099b1c210", + "customOutput": null, + "executionStartTime": 1659460454020, + "executionStopTime": 1659460454032 + }, + "source": [ + "our_structured.a = 21\n", + "print(MyConfigurable(**our_structured))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "eac7d385-9365-4098-acf9-4f0a0dbdcb85", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "eac7d385-9365-4098-acf9-4f0a0dbdcb85", + "customOutput": null, + "executionStartTime": 1659460599142, + "executionStopTime": 1659460599149 + }, + "source": [ + "## 4. First enhancement: nested types 🪺\n", + "\n", + "Our system allows Configurable classes to contain each other. \n", + "One thing to remember: add a call to `run_auto_creation` in `__post_init__`." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "9bd70ee5-4ec1-4021-bce5-9638b5088c0a", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "9bd70ee5-4ec1-4021-bce5-9638b5088c0a", + "customOutput": null, + "executionStartTime": 1659465752418, + "executionStopTime": 1659465752976 + }, + "source": [ + "class Inner(Configurable):\n", + " a: int = 8\n", + " b: bool = True\n", + " c: Tuple[int, ...] = (2, 3, 4, 6)\n", + "\n", + "\n", + "class Outer(Configurable):\n", + " inner: Inner\n", + " x: str = \"hello\"\n", + " xx: bool = False\n", + "\n", + " def __post_init__(self):\n", + " run_auto_creation(self)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "9f2b9f98-b54b-46cc-9b02-9e902cb279e7", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "9f2b9f98-b54b-46cc-9b02-9e902cb279e7", + "customOutput": null, + "executionStartTime": 1659465762326, + "executionStopTime": 1659465762339 + }, + "source": [ + "outer_dc = get_default_args(Outer)\n", + "print(OmegaConf.to_yaml(outer_dc))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "0254204b-8c7a-4d40-bba6-5132185f63d7", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "0254204b-8c7a-4d40-bba6-5132185f63d7", + "customOutput": null, + "executionStartTime": 1659465772894, + "executionStopTime": 1659465772911 + }, + "source": [ + "outer = Outer(**outer_dc)\n", + "assert isinstance(outer, Outer)\n", + "assert isinstance(outer.inner, Inner)\n", + "print(vars(outer))\n", + "print(outer.inner)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "44a78c13-ec92-4a87-808a-c4674b320c22", + "showInput": false, + "customInput": null + }, + "source": [ + "Note how inner_args is an extra member of outer. `run_auto_creation(self)` is equivalent to\n", + "```\n", + " self.inner = Inner(**self.inner_args)\n", + "```" + ], + "attachments": {} + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "af0ec78b-7888-4b0d-9346-63d970d43293", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "af0ec78b-7888-4b0d-9346-63d970d43293", + "customOutput": null, + "executionStartTime": 1659461071129, + "executionStopTime": 1659461071137 + }, + "source": [ + "## 5. Second enhancement: pluggable/replaceable components 🔌\n", + "\n", + "If a class uses `ReplaceableBase` as a base class instead of `Configurable`, we call it a replaceable.\n", + "It indicates that it is designed for child classes to use in its place.\n", + "We might use `NotImplementedError` to indicate functionality which subclasses are expected to implement.\n", + "The system maintains a global `registry` containing subclasses of each ReplaceableBase.\n", + "The subclasses register themselves with it with a decorator.\n", + "\n", + "A configurable class (i.e. a class which uses our system, i.e. a child of `Configurable` or `ReplaceableBase`) which contains a ReplaceableBase must also \n", + "contain a corresponding class_type field of type `str` which indicates which concrete child class to use." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "f2898703-d147-4394-978e-fc7f1f559395", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "f2898703-d147-4394-978e-fc7f1f559395", + "customOutput": null, + "executionStartTime": 1659463453457, + "executionStopTime": 1659463453467 + }, + "source": [ + "class InnerBase(ReplaceableBase):\n", + " def say_something(self):\n", + " raise NotImplementedError\n", + "\n", + "\n", + "@registry.register\n", + "class Inner1(InnerBase):\n", + " a: int = 1\n", + " b: str = \"h\"\n", + "\n", + " def say_something(self):\n", + " print(\"hello from an Inner1\")\n", + "\n", + "\n", + "@registry.register\n", + "class Inner2(InnerBase):\n", + " a: int = 2\n", + "\n", + " def say_something(self):\n", + " print(\"hello from an Inner2\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "6f171599-51ee-440f-82d7-a59f84d24624", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "6f171599-51ee-440f-82d7-a59f84d24624", + "customOutput": null, + "executionStartTime": 1659463453514, + "executionStopTime": 1659463453592 + }, + "source": [ + "class Out(Configurable):\n", + " inner: InnerBase\n", + " inner_class_type: str = \"Inner1\"\n", + " x: int = 19\n", + "\n", + " def __post_init__(self):\n", + " run_auto_creation(self)\n", + "\n", + " def talk(self):\n", + " self.inner.say_something()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "7abaecec-96e6-44df-8c8d-69c36a14b913", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "7abaecec-96e6-44df-8c8d-69c36a14b913", + "customOutput": null, + "executionStartTime": 1659463191360, + "executionStopTime": 1659463191428 + }, + "source": [ + "Out_dc = get_default_args(Out)\n", + "print(OmegaConf.to_yaml(Out_dc))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "c82dc2ca-ba8f-4a44-aed3-43f6b52ec28c", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "c82dc2ca-ba8f-4a44-aed3-43f6b52ec28c", + "customOutput": null, + "executionStartTime": 1659463192717, + "executionStopTime": 1659463192754 + }, + "source": [ + "Out_dc.inner_class_type = \"Inner2\"\n", + "out = Out(**Out_dc)\n", + "print(out.inner)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "aa0e1b04-963a-4724-81b7-5748b598b541", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "aa0e1b04-963a-4724-81b7-5748b598b541", + "customOutput": null, + "executionStartTime": 1659463193751, + "executionStopTime": 1659463193791 + }, + "source": [ + "out.talk()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "4f78a56c-39cd-4563-a97e-041e5f360f6b", + "showInput": false, + "customInput": null + }, + "source": [ + "Note in this case there are many `args` members. It is usually fine to ignore them in the code. They are needed for the config." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "ce7069d5-a813-4286-a7cd-6ff40362105a", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "ce7069d5-a813-4286-a7cd-6ff40362105a", + "customOutput": null, + "executionStartTime": 1659462145294, + "executionStopTime": 1659462145307 + }, + "source": [ + "print(vars(out))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "c7f051ff-c264-4b89-80dc-36cf179aafaf", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "c7f051ff-c264-4b89-80dc-36cf179aafaf", + "customOutput": null, + "executionStartTime": 1659462231114, + "executionStopTime": 1659462231130 + }, + "source": [ + "## 6. Example with torch.nn.Module 🔥\n", + "Typically in implicitron, we use this system in combination with [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)s. \n", + "Note in this case it is necessary to call `Module.__init__` explicitly in `__post_init__`." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "42d210d6-09e0-4daf-8ccb-411d30f268f4", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "42d210d6-09e0-4daf-8ccb-411d30f268f4", + "customOutput": null, + "executionStartTime": 1659462645018, + "executionStopTime": 1659462645037 + }, + "source": [ + "class MyLinear(torch.nn.Module, Configurable):\n", + " d_in: int = 2\n", + " d_out: int = 200\n", + "\n", + " def __post_init__(self):\n", + " super().__init__()\n", + " self.linear = torch.nn.Linear(in_features=self.d_in, out_features=self.d_out)\n", + "\n", + " def forward(self, x):\n", + " return self.linear.forward(x)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "546781fe-5b95-4e48-9cb5-34a634a31313", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "546781fe-5b95-4e48-9cb5-34a634a31313", + "customOutput": null, + "executionStartTime": 1659462692309, + "executionStopTime": 1659462692346 + }, + "source": [ + "expand_args_fields(MyLinear)\n", + "my_linear = MyLinear()\n", + "input = torch.zeros(2)\n", + "output = my_linear(input)\n", + "print(\"output shape:\", output.shape)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "b6cb71e1-1d54-4e89-a422-0a70772c5c03", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "b6cb71e1-1d54-4e89-a422-0a70772c5c03", + "customOutput": null, + "executionStartTime": 1659462738302, + "executionStopTime": 1659462738419 + }, + "source": [ + "`my_linear` has all the usual features of a Module.\n", + "E.g. it can be saved and loaded with `torch.save` and `torch.load`.\n", + "It has parameters:" + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "47e8c53e-2d2c-4b41-8aa3-65aa3ea8a7d3", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "47e8c53e-2d2c-4b41-8aa3-65aa3ea8a7d3", + "customOutput": null, + "executionStartTime": 1659462821485, + "executionStopTime": 1659462821501 + }, + "source": [ + "for name, value in my_linear.named_parameters():\n", + " print(name, value.shape)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "a01f0ea7-55f2-4af9-8e81-45dddf40f13b", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "a01f0ea7-55f2-4af9-8e81-45dddf40f13b", + "customOutput": null, + "executionStartTime": 1659463222379, + "executionStopTime": 1659463222409 + }, + "source": [ + "## 7. Example of implementing your own pluggable component \n", + "Let's say I am using a library with `Out` like in section **5** but I want to implement my own child of InnerBase. \n", + "All I need to do is register its definition, but I need to do this before expand_args_fields is explicitly or implicitly called on Out." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "d9635511-a52b-43d5-8dae-d5c1a3dd9157", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "d9635511-a52b-43d5-8dae-d5c1a3dd9157", + "customOutput": null, + "executionStartTime": 1659463694644, + "executionStopTime": 1659463694653 + }, + "source": [ + "@registry.register\n", + "class UserImplementedInner(InnerBase):\n", + " a: int = 200\n", + "\n", + " def say_something(self):\n", + " print(\"hello from the user\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "f1511aa2-56b8-4ed0-a453-17e2bbfeefe7", + "showInput": false, + "customInput": null + }, + "source": [ + "At this point, we need to redefine the class Out. \n", + "Otherwise if it has already been expanded without UserImplementedInner, then the following would not work,\n", + "because the implementations known to a class are fixed when it is expanded.\n", + "\n", + "If you are running experiments from a script, the thing to remember here is that you must import your own modules, which register your own implementations,\n", + "before you *use* the library classes." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "c7bb5a6e-682b-4eb0-a214-e0f5990b9406", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "c7bb5a6e-682b-4eb0-a214-e0f5990b9406", + "customOutput": null, + "executionStartTime": 1659463745967, + "executionStopTime": 1659463745986 + }, + "source": [ + "class Out(Configurable):\n", + " inner: InnerBase\n", + " inner_class_type: str = \"Inner1\"\n", + " x: int = 19\n", + "\n", + " def __post_init__(self):\n", + " run_auto_creation(self)\n", + "\n", + " def talk(self):\n", + " self.inner.say_something()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "b6ecdc86-4b7b-47c6-9f45-a7e557c94979", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "b6ecdc86-4b7b-47c6-9f45-a7e557c94979", + "customOutput": null, + "executionStartTime": 1659463747398, + "executionStopTime": 1659463747431 + }, + "source": [ + "expand_args_fields(Out)\n", + "out2 = Out(inner_class_type=\"UserImplementedInner\")\n", + "print(out2.inner)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "c7fe0df3-da13-40b8-9b06-6b1f37f37bb9", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "c7fe0df3-da13-40b8-9b06-6b1f37f37bb9", + "customOutput": null, + "executionStartTime": 1659464033633, + "executionStopTime": 1659464033643 + }, + "source": [ + "## 8: Example of making a subcomponent pluggable\n", + "\n", + "Let's look what needs to happen if we have a subcomponent which we make pluggable, to allow users to supply their own." + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "e37227b2-6897-4033-8560-9f2040abdeeb", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "e37227b2-6897-4033-8560-9f2040abdeeb", + "customOutput": null, + "executionStartTime": 1659464709922, + "executionStopTime": 1659464709933 + }, + "source": [ + "class SubComponent(Configurable):\n", + " x: float = 0.25\n", + "\n", + " def apply(self, a: float) -> float:\n", + " return a + self.x\n", + "\n", + "\n", + "class LargeComponent(Configurable):\n", + " repeats: int = 4\n", + " subcomponent: SubComponent\n", + "\n", + " def __post_init__(self):\n", + " run_auto_creation(self)\n", + "\n", + " def apply(self, a: float) -> float:\n", + " for _ in range(self.repeats):\n", + " a = self.subcomponent.apply(a)\n", + " return a" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "cab4c121-350e-443f-9a49-bd542a9735a2", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "cab4c121-350e-443f-9a49-bd542a9735a2", + "customOutput": null, + "executionStartTime": 1659464710339, + "executionStopTime": 1659464710459 + }, + "source": [ + "expand_args_fields(LargeComponent)\n", + "large_component = LargeComponent()\n", + "assert large_component.apply(3) == 4\n", + "print(OmegaConf.to_yaml(LargeComponent))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "be60323a-badf-46e4-a259-72cae1391028", + "showInput": false, + "customInput": null + }, + "source": [ + "Made generic:" + ], + "attachments": {} + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "fc0d8cdb-4627-4427-b92a-17ac1c1b37b8", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "fc0d8cdb-4627-4427-b92a-17ac1c1b37b8", + "customOutput": null, + "executionStartTime": 1659464717226, + "executionStopTime": 1659464717261 + }, + "source": [ + "class SubComponentBase(ReplaceableBase):\n", + " def apply(self, a: float) -> float:\n", + " raise NotImplementedError\n", + "\n", + "\n", + "@registry.register\n", + "class SubComponent(SubComponentBase):\n", + " x: float = 0.25\n", + "\n", + " def apply(self, a: float) -> float:\n", + " return a + self.x\n", + "\n", + "\n", + "class LargeComponent(Configurable):\n", + " repeats: int = 4\n", + " subcomponent: SubComponentBase\n", + " subcomponent_class_type: str = \"SubComponent\"\n", + "\n", + " def __post_init__(self):\n", + " run_auto_creation(self)\n", + "\n", + " def apply(self, a: float) -> float:\n", + " for _ in range(self.repeats):\n", + " a = self.subcomponent.apply(a)\n", + " return a" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "originalKey": "bbc3d321-6b49-4356-be75-1a173b1fc3a5", + "showInput": true, + "customInput": null, + "collapsed": false, + "requestMsgId": "bbc3d321-6b49-4356-be75-1a173b1fc3a5", + "customOutput": null, + "executionStartTime": 1659464725473, + "executionStopTime": 1659464725587 + }, + "source": [ + "expand_args_fields(LargeComponent)\n", + "large_component = LargeComponent()\n", + "assert large_component.apply(3) == 4\n", + "print(OmegaConf.to_yaml(LargeComponent))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "5115453a-1d96-4022-97e7-46433e6dcf60", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "5115453a-1d96-4022-97e7-46433e6dcf60", + "customOutput": null, + "executionStartTime": 1659464672680, + "executionStopTime": 1659464673231 + }, + "source": [ + "The following things had to change:\n", + "* The base class SubComponentBase was defined.\n", + "* SubComponent gained a `@registry.register` decoration and had its base class changed to the new one.\n", + "* `subcomponent_class_type` was added as a member of the outer class.\n", + "* In any saved configuration yaml files, the key `subcomponent_args` had to be changed to `subcomponent_SubComponent_args`." + ], + "attachments": {} + }, + { + "cell_type": "markdown", + "metadata": { + "originalKey": "0739269e-5c0e-4551-b06f-f4aab386ba54", + "showInput": false, + "customInput": null, + "collapsed": false, + "requestMsgId": "0739269e-5c0e-4551-b06f-f4aab386ba54", + "customOutput": null, + "executionStartTime": 1659462041307, + "executionStopTime": 1659462041637 + }, + "source": [ + "## Appendix: gotchas ⚠️\n", + "\n", + "* Omitting to define `__post_init__` or not calling `run_auto_creation` in it.\n", + "* Using a configurable class without calling get_default_args or expand_args_fields on it.\n", + "* Omitting a type annotation on a field. For example, writing \n", + "```\n", + " subcomponent_class_type = \"SubComponent\"\n", + "```\n", + "instead of \n", + "```\n", + " subcomponent_class_type: str = \"SubComponent\"\n", + "```\n", + "\n", + "" + ], + "attachments": {} + } + ] +} diff --git a/files/implicitron_config_system.py b/files/implicitron_config_system.py new file mode 100644 index 00000000..5904e8d5 --- /dev/null +++ b/files/implicitron_config_system.py @@ -0,0 +1,519 @@ + +# coding: utf-8 + +# In[ ]: + + +# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + + +# # Implicitron's config system + +# Implicitron's components are all based on a unified hierarchical configuration system. +# This allows configurable variables and all defaults to be defined separately for each new component. +# All configs relevant to an experiment are then automatically composed into a single configuration file that fully specifies the experiment. +# An especially important feature is extension points where users can insert their own sub-classes of Implicitron's base components. +# +# The file which defines this system is [here](https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/implicitron/tools/config.py) in the PyTorch3D repo. +# The Implicitron volumes tutorial contains a simple example of using the config system. +# This tutorial provides detailed hands-on experience in using and modifying Implicitron's configurable components. +# + +# ## 0. Install and import modules +# +# Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell: + +# In[ ]: + + +import os +import sys +import torch +need_pytorch3d=False +try: + import pytorch3d +except ModuleNotFoundError: + need_pytorch3d=True +if need_pytorch3d: + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): + # We try to install PyTorch3D via a released wheel. + pyt_version_str=torch.__version__.split("+")[0].replace(".", "") + version_str="".join([ + f"py3{sys.version_info.minor}_cu", + torch.version.cuda.replace(".",""), + f"_pyt{pyt_version_str}" + ]) + get_ipython().system('pip install fvcore iopath') + get_ipython().system('pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html') + else: + # We try to install PyTorch3D from source. + get_ipython().system('curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz') + get_ipython().system('tar xzf 1.10.0.tar.gz') + os.environ["CUB_HOME"] = os.getcwd() + "/cub-1.10.0" + get_ipython().system("pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'") + + +# Ensure omegaconf is installed. If not, run this cell. (It should not be necessary to restart the runtime.) + +# In[ ]: + + +get_ipython().system('pip install omegaconf') + + +# In[ ]: + + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from omegaconf import DictConfig, OmegaConf +from pytorch3d.implicitron.tools.config import ( + Configurable, + ReplaceableBase, + expand_args_fields, + get_default_args, + registry, + run_auto_creation, +) + + +# ## 1. Introducing dataclasses +# +# [Type hints](https://docs.python.org/3/library/typing.html) give a taxonomy of types in Python. [Dataclasses](https://docs.python.org/3/library/dataclasses.html) let you create a class based on a list of members which have names, types and possibly default values. The `__init__` function is created automatically, and calls a `__post_init__` function if present as a final step. For example + +# In[ ]: + + +@dataclass +class MyDataclass: + a: int + b: int = 8 + c: Optional[Tuple[int, ...]] = None + + def __post_init__(self): + print(f"created with a = {self.a}") + self.d = 2 * self.b + + +# In[ ]: + + +my_dataclass_instance = MyDataclass(a=18) +assert my_dataclass_instance.d == 16 + + +# 👷 Note that the `dataclass` decorator here is function which modifies the definition of the class itself. +# It runs immediately after the definition. +# Our config system requires that implicitron library code contains classes whose modified versions need to be aware of user-defined implementations. +# Therefore we need the modification of the class to be delayed. We don't use a decorator. +# + +# ## 2. Introducing omegaconf and OmegaConf.structured +# +# The [omegaconf](https://github.com/omry/omegaconf/) library provides a DictConfig class which is like a `dict` with str keys, but with extra features for ease-of-use as a configuration system. + +# In[ ]: + + +dc = DictConfig({"a": 2, "b": True, "c": None, "d": "hello"}) +assert dc.a == dc["a"] == 2 + + +# OmegaConf has a serialization to and from yaml. The [Hydra](https://hydra.cc/) library relies on this for its configuration files. + +# In[ ]: + + +print(OmegaConf.to_yaml(dc)) +assert OmegaConf.create(OmegaConf.to_yaml(dc)) == dc + + +# OmegaConf.structured provides a DictConfig from a dataclass or instance of a dataclass. Unlike a normal DictConfig, it is type-checked and only known keys can be added. + +# In[ ]: + + +structured = OmegaConf.structured(MyDataclass) +assert isinstance(structured, DictConfig) +print(structured) +print() +print(OmegaConf.to_yaml(structured)) + + +# `structured` knows it is missing a value for `a`. + +# Such an object has members compatible with the dataclass, so an initialisation can be performed as follows. + +# In[ ]: + + +structured.a = 21 +my_dataclass_instance2 = MyDataclass(**structured) +print(my_dataclass_instance2) + + +# You can also call OmegaConf.structured on an instance. + +# In[ ]: + + +structured_from_instance = OmegaConf.structured(my_dataclass_instance) +my_dataclass_instance3 = MyDataclass(**structured_from_instance) +print(my_dataclass_instance3) + + +# ## 3. Our approach to OmegaConf.structured +# +# We provide functions which are equivalent to `OmegaConf.structured` but support more features. +# To achieve the above using our functions, the following is used. +# Note that we indicate configurable classes using a special base class `Configurable`, not a decorator. + +# In[ ]: + + +class MyConfigurable(Configurable): + a: int + b: int = 8 + c: Optional[Tuple[int, ...]] = None + + def __post_init__(self): + print(f"created with a = {self.a}") + self.d = 2 * self.b + + +# In[ ]: + + +# expand_args_fields must be called on an object before it is instantiated. +# A warning is raised if this is missed, but it is possible to not notice the warning. +# It modifies the class like @dataclass +expand_args_fields(MyConfigurable) +my_configurable_instance = MyConfigurable(a=18) +assert my_configurable_instance.d == 16 + + +# In[ ]: + + +# get_default_args calls expand_args_fields automatically +our_structured = get_default_args(MyConfigurable) +assert isinstance(our_structured, DictConfig) +print(OmegaConf.to_yaml(our_structured)) + + +# In[ ]: + + +our_structured.a = 21 +print(MyConfigurable(**our_structured)) + + +# ## 4. First enhancement: nested types 🪺 +# +# Our system allows Configurable classes to contain each other. +# One thing to remember: add a call to `run_auto_creation` in `__post_init__`. + +# In[ ]: + + +class Inner(Configurable): + a: int = 8 + b: bool = True + c: Tuple[int, ...] = (2, 3, 4, 6) + + +class Outer(Configurable): + inner: Inner + x: str = "hello" + xx: bool = False + + def __post_init__(self): + run_auto_creation(self) + + +# In[ ]: + + +outer_dc = get_default_args(Outer) +print(OmegaConf.to_yaml(outer_dc)) + + +# In[ ]: + + +outer = Outer(**outer_dc) +assert isinstance(outer, Outer) +assert isinstance(outer.inner, Inner) +print(vars(outer)) +print(outer.inner) + + +# Note how inner_args is an extra member of outer. `run_auto_creation(self)` is equivalent to +# ``` +# self.inner = Inner(**self.inner_args) +# ``` + +# ## 5. Second enhancement: pluggable/replaceable components 🔌 +# +# If a class uses `ReplaceableBase` as a base class instead of `Configurable`, we call it a replaceable. +# It indicates that it is designed for child classes to use in its place. +# We might use `NotImplementedError` to indicate functionality which subclasses are expected to implement. +# The system maintains a global `registry` containing subclasses of each ReplaceableBase. +# The subclasses register themselves with it with a decorator. +# +# A configurable class (i.e. a class which uses our system, i.e. a child of `Configurable` or `ReplaceableBase`) which contains a ReplaceableBase must also +# contain a corresponding class_type field of type `str` which indicates which concrete child class to use. + +# In[ ]: + + +class InnerBase(ReplaceableBase): + def say_something(self): + raise NotImplementedError + + +@registry.register +class Inner1(InnerBase): + a: int = 1 + b: str = "h" + + def say_something(self): + print("hello from an Inner1") + + +@registry.register +class Inner2(InnerBase): + a: int = 2 + + def say_something(self): + print("hello from an Inner2") + + +# In[ ]: + + +class Out(Configurable): + inner: InnerBase + inner_class_type: str = "Inner1" + x: int = 19 + + def __post_init__(self): + run_auto_creation(self) + + def talk(self): + self.inner.say_something() + + +# In[ ]: + + +Out_dc = get_default_args(Out) +print(OmegaConf.to_yaml(Out_dc)) + + +# In[ ]: + + +Out_dc.inner_class_type = "Inner2" +out = Out(**Out_dc) +print(out.inner) + + +# In[ ]: + + +out.talk() + + +# Note in this case there are many `args` members. It is usually fine to ignore them in the code. They are needed for the config. + +# In[ ]: + + +print(vars(out)) + + +# ## 6. Example with torch.nn.Module 🔥 +# Typically in implicitron, we use this system in combination with [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)s. +# Note in this case it is necessary to call `Module.__init__` explicitly in `__post_init__`. + +# In[ ]: + + +class MyLinear(torch.nn.Module, Configurable): + d_in: int = 2 + d_out: int = 200 + + def __post_init__(self): + super().__init__() + self.linear = torch.nn.Linear(in_features=self.d_in, out_features=self.d_out) + + def forward(self, x): + return self.linear.forward(x) + + +# In[ ]: + + +expand_args_fields(MyLinear) +my_linear = MyLinear() +input = torch.zeros(2) +output = my_linear(input) +print("output shape:", output.shape) + + +# `my_linear` has all the usual features of a Module. +# E.g. it can be saved and loaded with `torch.save` and `torch.load`. +# It has parameters: + +# In[ ]: + + +for name, value in my_linear.named_parameters(): + print(name, value.shape) + + +# ## 7. Example of implementing your own pluggable component +# Let's say I am using a library with `Out` like in section **5** but I want to implement my own child of InnerBase. +# All I need to do is register its definition, but I need to do this before expand_args_fields is explicitly or implicitly called on Out. + +# In[ ]: + + +@registry.register +class UserImplementedInner(InnerBase): + a: int = 200 + + def say_something(self): + print("hello from the user") + + +# At this point, we need to redefine the class Out. +# Otherwise if it has already been expanded without UserImplementedInner, then the following would not work, +# because the implementations known to a class are fixed when it is expanded. +# +# If you are running experiments from a script, the thing to remember here is that you must import your own modules, which register your own implementations, +# before you *use* the library classes. + +# In[ ]: + + +class Out(Configurable): + inner: InnerBase + inner_class_type: str = "Inner1" + x: int = 19 + + def __post_init__(self): + run_auto_creation(self) + + def talk(self): + self.inner.say_something() + + +# In[ ]: + + +expand_args_fields(Out) +out2 = Out(inner_class_type="UserImplementedInner") +print(out2.inner) + + +# ## 8: Example of making a subcomponent pluggable +# +# Let's look what needs to happen if we have a subcomponent which we make pluggable, to allow users to supply their own. + +# In[ ]: + + +class SubComponent(Configurable): + x: float = 0.25 + + def apply(self, a: float) -> float: + return a + self.x + + +class LargeComponent(Configurable): + repeats: int = 4 + subcomponent: SubComponent + + def __post_init__(self): + run_auto_creation(self) + + def apply(self, a: float) -> float: + for _ in range(self.repeats): + a = self.subcomponent.apply(a) + return a + + +# In[ ]: + + +expand_args_fields(LargeComponent) +large_component = LargeComponent() +assert large_component.apply(3) == 4 +print(OmegaConf.to_yaml(LargeComponent)) + + +# Made generic: + +# In[ ]: + + +class SubComponentBase(ReplaceableBase): + def apply(self, a: float) -> float: + raise NotImplementedError + + +@registry.register +class SubComponent(SubComponentBase): + x: float = 0.25 + + def apply(self, a: float) -> float: + return a + self.x + + +class LargeComponent(Configurable): + repeats: int = 4 + subcomponent: SubComponentBase + subcomponent_class_type: str = "SubComponent" + + def __post_init__(self): + run_auto_creation(self) + + def apply(self, a: float) -> float: + for _ in range(self.repeats): + a = self.subcomponent.apply(a) + return a + + +# In[ ]: + + +expand_args_fields(LargeComponent) +large_component = LargeComponent() +assert large_component.apply(3) == 4 +print(OmegaConf.to_yaml(LargeComponent)) + + +# The following things had to change: +# * The base class SubComponentBase was defined. +# * SubComponent gained a `@registry.register` decoration and had its base class changed to the new one. +# * `subcomponent_class_type` was added as a member of the outer class. +# * In any saved configuration yaml files, the key `subcomponent_args` had to be changed to `subcomponent_SubComponent_args`. + +# ## Appendix: gotchas ⚠️ +# +# * Omitting to define `__post_init__` or not calling `run_auto_creation` in it. +# * Using a configurable class without calling get_default_args or expand_args_fields on it. +# * Omitting a type annotation on a field. For example, writing +# ``` +# subcomponent_class_type = "SubComponent" +# ``` +# instead of +# ``` +# subcomponent_class_type: str = "SubComponent" +# ``` +# +# diff --git a/files/implicitron_volumes.ipynb b/files/implicitron_volumes.ipynb new file mode 100644 index 00000000..69a364d4 --- /dev/null +++ b/files/implicitron_volumes.ipynb @@ -0,0 +1,913 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659619824914, + "executionStopTime": 1659619825485, + "originalKey": "d38652e8-200a-413c-a36a-f4d349b78a9d", + "requestMsgId": "641de8aa-0e42-4446-9304-c160a2d226bf", + "showInput": true + }, + "outputs": [], + "source": [ + "# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "a48a9dcf-e80f-474b-a0c4-2c9a765b15c5", + "showInput": false + }, + "source": [ + "# A simple model using Implicitron\n", + "\n", + "In this demo, we use the VolumeRenderer from PyTorch3D as a custom implicit function in Implicitron. We will see\n", + "* some of the main objects in Implicitron\n", + "* how to plug in a custom part of a model" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "51337c0e-ad27-4b75-ad6a-737dca5d7b95", + "showInput": false + }, + "source": [ + "## 0. Install and import modules\n", + "\n", + "Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659619898147, + "executionStopTime": 1659619898274, + "originalKey": "76f1ecd4-6b73-4214-81b0-118ef8d86872", + "requestMsgId": "deb6a860-6923-4227-abef-d31388b5142d", + "showInput": true + }, + "outputs": [], + "source": [ + "import os\n", + "import sys\n", + "import torch\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.12.\") 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", + " !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'" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "2c1020e6-eb4a-4644-9719-9147500d8e4f", + "showInput": false + }, + "source": [ + "Ensure omegaconf and visdom are installed. If not, run this cell. (It should not be necessary to restart the runtime.)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "customInput": null, + "customOutput": null, + "originalKey": "9e751931-a38d-44c9-9ff1-ac2f7d3a3f99", + "showInput": true + }, + "outputs": [], + "source": [ + "!pip install omegaconf visdom" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customOutput": null, + "executionStartTime": 1659612480556, + "executionStopTime": 1659612480644, + "hidden_ranges": [], + "originalKey": "86807e4a-1675-4520-a033-c7af85b233ec", + "requestMsgId": "880a7e20-4a90-4b37-a5eb-bccc0b23cac6" + }, + "outputs": [], + "source": [ + "import logging\n", + "from typing import Tuple\n", + "\n", + "import matplotlib.animation as animation\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import tqdm\n", + "from IPython.display import HTML\n", + "from omegaconf import OmegaConf\n", + "from PIL import Image\n", + "from pytorch3d.implicitron.dataset.dataset_base import FrameData\n", + "from pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider import RenderedMeshDatasetMapProvider\n", + "from pytorch3d.implicitron.models.generic_model import GenericModel\n", + "from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase\n", + "from pytorch3d.implicitron.models.renderer.base import EvaluationMode\n", + "from pytorch3d.implicitron.tools.config import expand_args_fields, get_default_args, registry, remove_unused_components\n", + "from pytorch3d.renderer import RayBundle\n", + "from pytorch3d.renderer.implicit.renderer import VolumeSampler\n", + "from pytorch3d.structures import Volumes\n", + "from pytorch3d.vis.plotly_vis import plot_batch_individually, plot_scene" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659610929375, + "executionStopTime": 1659610929383, + "hidden_ranges": [], + "originalKey": "b2d9f5bd-a9d4-4f78-b21e-92f2658e0fe9", + "requestMsgId": "7e43e623-4030-438b-af4e-b96170c9a052", + "showInput": true + }, + "outputs": [], + "source": [ + "output_resolution = 80" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659610930042, + "executionStopTime": 1659610930050, + "hidden_ranges": [], + "originalKey": "0b0c2087-4c86-4c57-b0ee-6f48a70a9c78", + "requestMsgId": "46883aad-f00b-4fd4-ac17-eec0b2ac272a", + "showInput": true + }, + "outputs": [], + "source": [ + "torch.set_printoptions(sci_mode=False)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "37809d0d-b02e-42df-85b6-cdd038373653", + "showInput": false + }, + "source": [ + "## 1. Load renders of a mesh (the cow mesh) as a dataset\n", + "\n", + "A dataset's train, val and test parts in Implicitron are represented as a `dataset_map`, and provided by an implementation of `DatasetMapProvider`. \n", + "`RenderedMeshDatasetMapProvider` is one which generates a single-scene dataset with only a train component by taking a mesh and rendering it.\n", + "We use it with the cow mesh." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659620739780, + "executionStopTime": 1659620739914, + "originalKey": "cc68cb9c-b8bf-4e9e-bef1-2cfafdf6caa2", + "requestMsgId": "398cfcae-5d43-4b6f-9c75-db3d297364d4", + "showInput": false + }, + "source": [ + "If running this notebook using **Google Colab**, run the following cell to fetch the mesh obj and texture files and save it at the path data/cow_mesh.\n", + "If running locally, the data is already available at the correct path." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "customInput": null, + "customOutput": null, + "originalKey": "2c55e002-a885-4169-8fdc-af9078b05968", + "showInput": true + }, + "outputs": [], + "source": [ + "!mkdir -p data/cow_mesh\n", + "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.obj\n", + "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.mtl\n", + "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow_texture.png" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "2a976be8-01bf-4a1c-a6e7-61d5d08c3dbd", + "showInput": false + }, + "source": [ + "If we want to instantiate one of Implicitron's configurable objects, such as `RenderedMeshDatasetMapProvider`, without using the OmegaConf initialisation (get_default_args), we need to call `expand_args_fields` on the class first." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customOutput": null, + "executionStartTime": 1659621652237, + "executionStopTime": 1659621652903, + "hidden_ranges": [], + "originalKey": "eb77aaec-048c-40bd-bd69-0e66b6ab60b1", + "requestMsgId": "09b9975c-ff86-41c9-b4a9-975d23afc562", + "showInput": true + }, + "outputs": [], + "source": [ + "expand_args_fields(RenderedMeshDatasetMapProvider)\n", + "cow_provider = RenderedMeshDatasetMapProvider(\n", + " data_file=\"data/cow_mesh/cow.obj\",\n", + " use_point_light=False,\n", + " resolution=output_resolution,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659610966145, + "executionStopTime": 1659610966255, + "hidden_ranges": [], + "originalKey": "8210e15b-da48-4306-a49a-41c4e7e7d42f", + "requestMsgId": "c243edd2-a106-4fba-8471-dfa4f99a2088", + "showInput": true + }, + "outputs": [], + "source": [ + "dataset_map = cow_provider.get_dataset_map()\n", + "tr_cameras = [training_frame.camera for training_frame in dataset_map.train]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659610967703, + "executionStopTime": 1659610967848, + "hidden_ranges": [], + "originalKey": "458d72ad-d9a7-4f13-b5b7-90d2aec61c16", + "requestMsgId": "7f9431f3-8717-4d89-a7fe-1420dd0e00c4", + "showInput": true + }, + "outputs": [], + "source": [ + "# The cameras are all in the XZ plane, in a circle about 2.7 from the origin\n", + "centers = torch.cat([i.get_camera_center() for i in tr_cameras])\n", + "print(centers.min(0).values)\n", + "print(centers.max(0).values)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659552920194, + "executionStopTime": 1659552923122, + "hidden_ranges": [], + "originalKey": "931e712b-b141-437a-97fb-dc2a07ce3458", + "requestMsgId": "931e712b-b141-437a-97fb-dc2a07ce3458", + "showInput": true + }, + "outputs": [], + "source": [ + "# visualization of the cameras\n", + "plot = plot_scene({\"k\": {i: camera for i, camera in enumerate(tr_cameras)}}, camera_scale=0.25)\n", + "plot.layout.scene.aspectmode = \"data\"\n", + "plot" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "afa9c02d-f76b-4f68-83e9-9733c615406b", + "showInput": false + }, + "source": [ + "## 2. Custom implicit function 🧊\n", + "\n", + "At the core of neural rendering methods are functions of spatial coordinates called implicit functions, which are used in some kind of rendering process.\n", + "(Often those functions can additionally take other data as well, such as view direction.)\n", + "A common rendering process is ray marching over densities and colors provided by an implicit function.\n", + "In our case, taking samples from a 3D volume grid is a very simple function of spatial coordinates. \n", + "\n", + "Here we define our own implicit function, which uses PyTorch3D's existing functionality for sampling from a volume grid.\n", + "We do this by subclassing `ImplicitFunctionBase`.\n", + "We need to register our subclass with a special decorator.\n", + "We use Python's dataclass annotations for configuring the module." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659613575850, + "executionStopTime": 1659613575940, + "hidden_ranges": [], + "originalKey": "61b55043-dc52-4de7-992e-e2195edd2123", + "requestMsgId": "dfaace3c-098c-4ffe-9240-6a7ae0ff271e", + "showInput": true + }, + "outputs": [], + "source": [ + "@registry.register\n", + "class MyVolumes(ImplicitFunctionBase, torch.nn.Module):\n", + " grid_resolution: int = 50 # common HWD of volumes, the number of voxels in each direction\n", + " extent: float = 1.0 # In world coordinates, the volume occupies is [-extent, extent] along each axis\n", + "\n", + " def __post_init__(self):\n", + " # We have to call this explicitly if there are other base classes like Module\n", + " super().__init__()\n", + "\n", + " # We define parameters like other torch.nn.Module objects.\n", + " # In this case, both our parameter tensors are trainable; they govern the contents of the volume grid.\n", + " density = torch.full((self.grid_resolution, self.grid_resolution, self.grid_resolution), -2.0)\n", + " self.density = torch.nn.Parameter(density)\n", + " color = torch.full((3, self.grid_resolution, self.grid_resolution, self.grid_resolution), 0.0)\n", + " self.color = torch.nn.Parameter(color)\n", + " self.density_activation = torch.nn.Softplus()\n", + "\n", + " def forward(\n", + " self,\n", + " ray_bundle: RayBundle,\n", + " fun_viewpool=None,\n", + " global_code=None,\n", + " ):\n", + " densities = self.density_activation(self.density[None, None])\n", + " voxel_size = 2.0 * float(self.extent) / self.grid_resolution\n", + " features = self.color.sigmoid()[None]\n", + "\n", + " # Like other PyTorch3D structures, the actual Volumes object should only exist as long\n", + " # as one iteration of training. It is local to this function.\n", + "\n", + " volume = Volumes(densities=densities, features=features, voxel_size=voxel_size)\n", + " sampler = VolumeSampler(volumes=volume)\n", + " densities, features = sampler(ray_bundle)\n", + "\n", + " # When an implicit function is used for raymarching, i.e. for MultiPassEmissionAbsorptionRenderer,\n", + " # it must return (densities, features, an auxiliary tuple)\n", + " return densities, features, {}\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "abaf2cd6-1b68-400e-a142-8fb9f49953f3", + "showInput": false + }, + "source": [ + "## 3. Construct the model object.\n", + "\n", + "The main model object in PyTorch3D is `GenericModel`, which has pluggable components for the major steps, including the renderer and the implicit function(s).\n", + "There are two ways to construct it which are equivalent here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621267561, + "executionStopTime": 1659621267938, + "originalKey": "f26c3dce-fbae-4592-bd0e-e4a8abc57c2c", + "requestMsgId": "9213687e-1caf-46a8-a4e5-a9c531530092", + "showInput": true + }, + "outputs": [], + "source": [ + "CONSTRUCT_MODEL_FROM_CONFIG = True\n", + "if CONSTRUCT_MODEL_FROM_CONFIG:\n", + " # Via a DictConfig - this is how our training loop with hydra works\n", + " cfg = get_default_args(GenericModel)\n", + " cfg.implicit_function_class_type = \"MyVolumes\"\n", + " cfg.render_image_height=output_resolution\n", + " cfg.render_image_width=output_resolution\n", + " cfg.loss_weights={\"loss_rgb_huber\": 1.0}\n", + " cfg.tqdm_trigger_threshold=19000\n", + " cfg.raysampler_AdaptiveRaySampler_args.scene_extent= 4.0\n", + " gm = GenericModel(**cfg)\n", + "else:\n", + " # constructing GenericModel directly\n", + " expand_args_fields(GenericModel)\n", + " gm = GenericModel(\n", + " implicit_function_class_type=\"MyVolumes\",\n", + " render_image_height=output_resolution,\n", + " render_image_width=output_resolution,\n", + " loss_weights={\"loss_rgb_huber\": 1.0},\n", + " tqdm_trigger_threshold=19000,\n", + " raysampler_AdaptiveRaySampler_args = {\"scene_extent\": 4.0}\n", + " )\n", + "\n", + " # In this case we can get the equivalent DictConfig cfg object to the way gm is configured as follows\n", + " cfg = OmegaConf.structured(gm)\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659611214689, + "executionStopTime": 1659611214748, + "hidden_ranges": [], + "originalKey": "4e659f7d-ce66-4999-83de-005eb09d7705", + "requestMsgId": "7b815b2b-cf19-44d0-ae89-76fde6df35ec", + "showInput": false + }, + "source": [ + " The default renderer is an emission-absorbtion raymarcher. We keep that default." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621268007, + "executionStopTime": 1659621268190, + "hidden_ranges": [], + "originalKey": "d37ae488-c57c-44d3-9def-825dc1a6495b", + "requestMsgId": "71143ec1-730f-4876-8a14-e46eea9d6dd1", + "showInput": true + }, + "outputs": [], + "source": [ + "# We can display the configuration in use as follows.\n", + "remove_unused_components(cfg)\n", + "yaml = OmegaConf.to_yaml(cfg, sort_keys=False)\n", + "%page -r yaml" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621268727, + "executionStopTime": 1659621268776, + "hidden_ranges": [], + "originalKey": "52e53179-3c6e-4c1f-a38a-3a6d803687bb", + "requestMsgId": "05de9bc3-3f74-4a6f-851c-9ec919b59506", + "showInput": true + }, + "outputs": [], + "source": [ + "device = torch.device(\"cuda:0\")\n", + "gm.to(device)\n", + "assert next(gm.parameters()).is_cuda" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "customInput": null, + "originalKey": "528a7d53-c645-49c2-9021-09adbb18cd23", + "showInput": false + }, + "source": [ + "## 4. train the model " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621270236, + "executionStopTime": 1659621270446, + "hidden_ranges": [], + "originalKey": "953280bd-3161-42ba-8dcb-0c8ef2d5cc25", + "requestMsgId": "9bba424b-7bfd-4e5a-9d79-ae316e20bab0", + "showInput": true + }, + "outputs": [], + "source": [ + "train_data_collated = [FrameData.collate([frame.to(device)]) for frame in dataset_map.train]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621270815, + "executionStopTime": 1659621270948, + "hidden_ranges": [], + "originalKey": "2fcf07f0-0c28-49c7-8c76-1c9a9d810167", + "requestMsgId": "821deb43-6084-4ece-83c3-dee214562c47", + "showInput": true + }, + "outputs": [], + "source": [ + "gm.train()\n", + "optimizer = torch.optim.Adam(gm.parameters(), lr=0.1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customOutput": null, + "executionStartTime": 1659621271875, + "executionStopTime": 1659621298146, + "hidden_ranges": [], + "originalKey": "105099f7-ed0c-4e7f-a976-61a93fd0a8fe", + "requestMsgId": "0c87c108-83e3-4129-ad02-85e0140f1368", + "showInput": true + }, + "outputs": [], + "source": [ + "iterator = tqdm.tqdm(range(2000))\n", + "for n_batch in iterator:\n", + " optimizer.zero_grad()\n", + "\n", + " frame = train_data_collated[n_batch % len(dataset_map.train)]\n", + " out = gm(**frame, evaluation_mode=EvaluationMode.TRAINING)\n", + " out[\"objective\"].backward()\n", + " if n_batch % 100 == 0:\n", + " iterator.set_postfix_str(f\"loss: {float(out['objective']):.5f}\")\n", + " optimizer.step()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659535024768, + "executionStopTime": 1659535024906, + "originalKey": "e3cd494a-536b-48bc-8290-c048118c82eb", + "requestMsgId": "e3cd494a-536b-48bc-8290-c048118c82eb", + "showInput": false + }, + "source": [ + "## 5. Evaluate the module\n", + "\n", + "We generate complete images from all the viewpoints to see how they look." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621299859, + "executionStopTime": 1659621311133, + "hidden_ranges": [], + "originalKey": "fbe1b2ea-cc24-4b20-a2d7-0249185e34a5", + "requestMsgId": "771ef1f8-5eee-4932-9e81-33604bf0512a", + "showInput": true + }, + "outputs": [], + "source": [ + "def to_numpy_image(image):\n", + " # Takes an image of shape (C, H, W) in [0,1], where C=3 or 1\n", + " # to a numpy uint image of shape (H, W, 3)\n", + " return (image * 255).to(torch.uint8).permute(1, 2, 0).detach().cpu().expand(-1, -1, 3).numpy()\n", + "def resize_image(image):\n", + " # Takes images of shape (B, C, H, W) to (B, C, output_resolution, output_resolution)\n", + " return torch.nn.functional.interpolate(image, size=(output_resolution, output_resolution))\n", + "\n", + "gm.eval()\n", + "images = []\n", + "expected = []\n", + "masks = []\n", + "masks_expected = []\n", + "for frame in tqdm.tqdm(train_data_collated):\n", + " with torch.no_grad():\n", + " out = gm(**frame, evaluation_mode=EvaluationMode.EVALUATION)\n", + "\n", + " image_rgb = to_numpy_image(out[\"images_render\"][0])\n", + " mask = to_numpy_image(out[\"masks_render\"][0])\n", + " expd = to_numpy_image(resize_image(frame.image_rgb)[0])\n", + " mask_expected = to_numpy_image(resize_image(frame.fg_probability)[0])\n", + "\n", + " images.append(image_rgb)\n", + " masks.append(mask)\n", + " expected.append(expd)\n", + " masks_expected.append(mask_expected)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659614622542, + "executionStopTime": 1659614622757, + "originalKey": "24953039-9780-40fd-bd81-5d63e9f40069", + "requestMsgId": "7af895a3-dfe4-4c28-ac3b-4ff0fbb40c7f", + "showInput": false + }, + "source": [ + "We draw a grid showing predicted image and expected image, followed by predicted mask and expected mask, from each viewpoint. \n", + "This is a grid of four rows of images, wrapped in to several large rows, i.e..\n", + "
\n", + "```\n", + "┌────────┬────────┐ ┌────────┐\n", + "│pred │pred │ │pred │\n", + "│image │image │ │image │\n", + "│1 │2 │ │n │\n", + "├────────┼────────┤ ├────────┤\n", + "│expected│expected│ │expected│\n", + "│image │image │ ... │image │\n", + "│1 │2 │ │n │\n", + "├────────┼────────┤ ├────────┤\n", + "│pred │pred │ │pred │\n", + "│mask │mask │ │mask │\n", + "│1 │2 │ │n │\n", + "├────────┼────────┤ ├────────┤\n", + "│expected│expected│ │expected│\n", + "│mask │mask │ │mask │\n", + "│1 │2 │ │n │\n", + "├────────┼────────┤ ├────────┤\n", + "│pred │pred │ │pred │\n", + "│image │image │ │image │\n", + "│n+1 │n+1 │ │2n │\n", + "├────────┼────────┤ ├────────┤\n", + "│expected│expected│ │expected│\n", + "│image │image │ ... │image │\n", + "│n+1 │n+2 │ │2n │\n", + "├────────┼────────┤ ├────────┤\n", + "│pred │pred │ │pred │\n", + "│mask │mask │ │mask │\n", + "│n+1 │n+2 │ │2n │\n", + "├────────┼────────┤ ├────────┤\n", + "│expected│expected│ │expected│\n", + "│mask │mask │ │mask │\n", + "│n+1 │n+2 │ │2n │\n", + "└────────┴────────┘ └────────┘\n", + " ...\n", + "```\n", + "
" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621313894, + "executionStopTime": 1659621314042, + "hidden_ranges": [], + "originalKey": "c488a34a-e46d-4649-93fb-4b1bb5a0e439", + "requestMsgId": "4221e632-fca1-4fe5-b2e3-f92c37aa40e4", + "showInput": true + }, + "outputs": [], + "source": [ + "images_to_display = [images.copy(), expected.copy(), masks.copy(), masks_expected.copy()]\n", + "n_rows = 4\n", + "n_images = len(images)\n", + "blank_image = images[0] * 0\n", + "n_per_row = 1+(n_images-1)//n_rows\n", + "for _ in range(n_per_row*n_rows - n_images):\n", + " for group in images_to_display:\n", + " group.append(blank_image)\n", + "\n", + "images_to_display_listed = [[[i] for i in j] for j in images_to_display]\n", + "split = []\n", + "for row in range(n_rows):\n", + " for group in images_to_display_listed:\n", + " split.append(group[row*n_per_row:(row+1)*n_per_row]) \n", + "\n", + "Image.fromarray(np.block(split))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621323795, + "executionStopTime": 1659621323820, + "hidden_ranges": [], + "originalKey": "49eab9e1-4fe2-4fbe-b4f3-7b6953340170", + "requestMsgId": "85b402ad-f903-431f-a13e-c2d697e869bb", + "showInput": true + }, + "outputs": [], + "source": [ + "# Print the maximum channel intensity in the first image.\n", + "print(images[1].max()/255)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "code_folding": [], + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621408642, + "executionStopTime": 1659621409559, + "hidden_ranges": [], + "originalKey": "137d2c43-d39d-4266-ac5e-2b714da5e0ee", + "requestMsgId": "8e27ec57-c2d6-4ae0-be69-b63b6af929ff", + "showInput": true + }, + "outputs": [], + "source": [ + "plt.ioff()\n", + "fig, ax = plt.subplots(figsize=(3,3))\n", + "\n", + "ax.grid(None)\n", + "ims = [[ax.imshow(im, animated=True)] for im in images]\n", + "ani = animation.ArtistAnimation(fig, ims, interval=80, blit=True)\n", + "ani_html = ani.to_jshtml()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659621409620, + "executionStopTime": 1659621409725, + "originalKey": "783e70d6-7cf1-4d76-a126-ba11ffc2f5be", + "requestMsgId": "b6843506-c5fa-4508-80fc-8ecae51a934a", + "showInput": true + }, + "outputs": [], + "source": [ + "HTML(ani_html)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "customInput": null, + "customOutput": null, + "executionStartTime": 1659614670081, + "executionStopTime": 1659614670168, + "originalKey": "0286c350-2362-4f47-8181-2fc2ba51cfcf", + "requestMsgId": "976f4db9-d4c7-466c-bcfd-218234400226", + "showInput": true + }, + "outputs": [], + "source": [ + "# If you want to see the output of the model with the volume forced to opaque white, run this and re-evaluate\n", + "# with torch.no_grad():\n", + "# gm._implicit_functions[0]._fn.density.fill_(9.0)\n", + "# gm._implicit_functions[0]._fn.color.fill_(9.0)\n" + ] + } + ], + "metadata": { + "bento_stylesheets": { + "bento/extensions/flow/main.css": true, + "bento/extensions/kernel_selector/main.css": true, + "bento/extensions/kernel_ui/main.css": true, + "bento/extensions/new_kernel/main.css": true, + "bento/extensions/system_usage/main.css": true, + "bento/extensions/theme/main.css": true + }, + "captumWidgetMessage": {}, + "dataExplorerConfig": {}, + "kernelspec": { + "display_name": "pytorch3d", + "language": "python", + "metadata": { + "cinder_runtime": false, + "fbpkg_supported": true, + "is_prebuilt": true, + "kernel_name": "bento_kernel_pytorch3d", + "nightly_builds": true + }, + "name": "bento_kernel_pytorch3d" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + }, + "last_base_url": "https://9177.od.fbinfra.net:443/", + "last_kernel_id": "bb33cd83-7924-489a-8bd8-2d9d62eb0126", + "last_msg_id": "99f7088e-d22b355b859660479ef0574e_5743", + "last_server_session_id": "2944b203-9ea8-4c0e-9634-645dfea5f26b", + "outputWidgetContext": {} + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/files/implicitron_volumes.py b/files/implicitron_volumes.py new file mode 100644 index 00000000..932b42c4 --- /dev/null +++ b/files/implicitron_volumes.py @@ -0,0 +1,420 @@ + +# coding: utf-8 + +# In[ ]: + + +# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + + +# # A simple model using Implicitron +# +# In this demo, we use the VolumeRenderer from PyTorch3D as a custom implicit function in Implicitron. We will see +# * some of the main objects in Implicitron +# * how to plug in a custom part of a model + +# ## 0. Install and import modules +# +# Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell: +# + +# In[ ]: + + +import os +import sys +import torch +need_pytorch3d=False +try: + import pytorch3d +except ModuleNotFoundError: + need_pytorch3d=True +if need_pytorch3d: + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): + # We try to install PyTorch3D via a released wheel. + pyt_version_str=torch.__version__.split("+")[0].replace(".", "") + version_str="".join([ + f"py3{sys.version_info.minor}_cu", + torch.version.cuda.replace(".",""), + f"_pyt{pyt_version_str}" + ]) + get_ipython().system('pip install fvcore iopath') + get_ipython().system('pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html') + else: + # We try to install PyTorch3D from source. + get_ipython().system('curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz') + get_ipython().system('tar xzf 1.10.0.tar.gz') + os.environ["CUB_HOME"] = os.getcwd() + "/cub-1.10.0" + get_ipython().system("pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'") + + +# Ensure omegaconf and visdom are installed. If not, run this cell. (It should not be necessary to restart the runtime.) + +# In[ ]: + + +get_ipython().system('pip install omegaconf visdom') + + +# In[ ]: + + +import logging +from typing import Tuple + +import matplotlib.animation as animation +import matplotlib.pyplot as plt +import numpy as np +import torch +import tqdm +from IPython.display import HTML +from omegaconf import OmegaConf +from PIL import Image +from pytorch3d.implicitron.dataset.dataset_base import FrameData +from pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider import RenderedMeshDatasetMapProvider +from pytorch3d.implicitron.models.generic_model import GenericModel +from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase +from pytorch3d.implicitron.models.renderer.base import EvaluationMode +from pytorch3d.implicitron.tools.config import expand_args_fields, get_default_args, registry, remove_unused_components +from pytorch3d.renderer import RayBundle +from pytorch3d.renderer.implicit.renderer import VolumeSampler +from pytorch3d.structures import Volumes +from pytorch3d.vis.plotly_vis import plot_batch_individually, plot_scene + + +# In[ ]: + + +output_resolution = 80 + + +# In[ ]: + + +torch.set_printoptions(sci_mode=False) + + +# ## 1. Load renders of a mesh (the cow mesh) as a dataset +# +# A dataset's train, val and test parts in Implicitron are represented as a `dataset_map`, and provided by an implementation of `DatasetMapProvider`. +# `RenderedMeshDatasetMapProvider` is one which generates a single-scene dataset with only a train component by taking a mesh and rendering it. +# We use it with the cow mesh. + +# If running this notebook using **Google Colab**, run the following cell to fetch the mesh obj and texture files and save it at the path data/cow_mesh. +# If running locally, the data is already available at the correct path. + +# In[ ]: + + +get_ipython().system('mkdir -p data/cow_mesh') +get_ipython().system('wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.obj') +get_ipython().system('wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.mtl') +get_ipython().system('wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow_texture.png') + + +# If we want to instantiate one of Implicitron's configurable objects, such as `RenderedMeshDatasetMapProvider`, without using the OmegaConf initialisation (get_default_args), we need to call `expand_args_fields` on the class first. + +# In[ ]: + + +expand_args_fields(RenderedMeshDatasetMapProvider) +cow_provider = RenderedMeshDatasetMapProvider( + data_file="data/cow_mesh/cow.obj", + use_point_light=False, + resolution=output_resolution, +) + + +# In[ ]: + + +dataset_map = cow_provider.get_dataset_map() +tr_cameras = [training_frame.camera for training_frame in dataset_map.train] + + +# In[ ]: + + +# The cameras are all in the XZ plane, in a circle about 2.7 from the origin +centers = torch.cat([i.get_camera_center() for i in tr_cameras]) +print(centers.min(0).values) +print(centers.max(0).values) + + +# In[ ]: + + +# visualization of the cameras +plot = plot_scene({"k": {i: camera for i, camera in enumerate(tr_cameras)}}, camera_scale=0.25) +plot.layout.scene.aspectmode = "data" +plot + + +# ## 2. Custom implicit function 🧊 +# +# At the core of neural rendering methods are functions of spatial coordinates called implicit functions, which are used in some kind of rendering process. +# (Often those functions can additionally take other data as well, such as view direction.) +# A common rendering process is ray marching over densities and colors provided by an implicit function. +# In our case, taking samples from a 3D volume grid is a very simple function of spatial coordinates. +# +# Here we define our own implicit function, which uses PyTorch3D's existing functionality for sampling from a volume grid. +# We do this by subclassing `ImplicitFunctionBase`. +# We need to register our subclass with a special decorator. +# We use Python's dataclass annotations for configuring the module. + +# In[ ]: + + +@registry.register +class MyVolumes(ImplicitFunctionBase, torch.nn.Module): + grid_resolution: int = 50 # common HWD of volumes, the number of voxels in each direction + extent: float = 1.0 # In world coordinates, the volume occupies is [-extent, extent] along each axis + + def __post_init__(self): + # We have to call this explicitly if there are other base classes like Module + super().__init__() + + # We define parameters like other torch.nn.Module objects. + # In this case, both our parameter tensors are trainable; they govern the contents of the volume grid. + density = torch.full((self.grid_resolution, self.grid_resolution, self.grid_resolution), -2.0) + self.density = torch.nn.Parameter(density) + color = torch.full((3, self.grid_resolution, self.grid_resolution, self.grid_resolution), 0.0) + self.color = torch.nn.Parameter(color) + self.density_activation = torch.nn.Softplus() + + def forward( + self, + ray_bundle: RayBundle, + fun_viewpool=None, + global_code=None, + ): + densities = self.density_activation(self.density[None, None]) + voxel_size = 2.0 * float(self.extent) / self.grid_resolution + features = self.color.sigmoid()[None] + + # Like other PyTorch3D structures, the actual Volumes object should only exist as long + # as one iteration of training. It is local to this function. + + volume = Volumes(densities=densities, features=features, voxel_size=voxel_size) + sampler = VolumeSampler(volumes=volume) + densities, features = sampler(ray_bundle) + + # When an implicit function is used for raymarching, i.e. for MultiPassEmissionAbsorptionRenderer, + # it must return (densities, features, an auxiliary tuple) + return densities, features, {} + + +# ## 3. Construct the model object. +# +# The main model object in PyTorch3D is `GenericModel`, which has pluggable components for the major steps, including the renderer and the implicit function(s). +# There are two ways to construct it which are equivalent here. + +# In[ ]: + + +CONSTRUCT_MODEL_FROM_CONFIG = True +if CONSTRUCT_MODEL_FROM_CONFIG: + # Via a DictConfig - this is how our training loop with hydra works + cfg = get_default_args(GenericModel) + cfg.implicit_function_class_type = "MyVolumes" + cfg.render_image_height=output_resolution + cfg.render_image_width=output_resolution + cfg.loss_weights={"loss_rgb_huber": 1.0} + cfg.tqdm_trigger_threshold=19000 + cfg.raysampler_AdaptiveRaySampler_args.scene_extent= 4.0 + gm = GenericModel(**cfg) +else: + # constructing GenericModel directly + expand_args_fields(GenericModel) + gm = GenericModel( + implicit_function_class_type="MyVolumes", + render_image_height=output_resolution, + render_image_width=output_resolution, + loss_weights={"loss_rgb_huber": 1.0}, + tqdm_trigger_threshold=19000, + raysampler_AdaptiveRaySampler_args = {"scene_extent": 4.0} + ) + + # In this case we can get the equivalent DictConfig cfg object to the way gm is configured as follows + cfg = OmegaConf.structured(gm) + + +# The default renderer is an emission-absorbtion raymarcher. We keep that default. + +# In[ ]: + + +# We can display the configuration in use as follows. +remove_unused_components(cfg) +yaml = OmegaConf.to_yaml(cfg, sort_keys=False) +get_ipython().run_line_magic('page', '-r yaml') + + +# In[ ]: + + +device = torch.device("cuda:0") +gm.to(device) +assert next(gm.parameters()).is_cuda + + +# ## 4. train the model + +# In[ ]: + + +train_data_collated = [FrameData.collate([frame.to(device)]) for frame in dataset_map.train] + + +# In[ ]: + + +gm.train() +optimizer = torch.optim.Adam(gm.parameters(), lr=0.1) + + +# In[ ]: + + +iterator = tqdm.tqdm(range(2000)) +for n_batch in iterator: + optimizer.zero_grad() + + frame = train_data_collated[n_batch % len(dataset_map.train)] + out = gm(**frame, evaluation_mode=EvaluationMode.TRAINING) + out["objective"].backward() + if n_batch % 100 == 0: + iterator.set_postfix_str(f"loss: {float(out['objective']):.5f}") + optimizer.step() + + +# ## 5. Evaluate the module +# +# We generate complete images from all the viewpoints to see how they look. + +# In[ ]: + + +def to_numpy_image(image): + # Takes an image of shape (C, H, W) in [0,1], where C=3 or 1 + # to a numpy uint image of shape (H, W, 3) + return (image * 255).to(torch.uint8).permute(1, 2, 0).detach().cpu().expand(-1, -1, 3).numpy() +def resize_image(image): + # Takes images of shape (B, C, H, W) to (B, C, output_resolution, output_resolution) + return torch.nn.functional.interpolate(image, size=(output_resolution, output_resolution)) + +gm.eval() +images = [] +expected = [] +masks = [] +masks_expected = [] +for frame in tqdm.tqdm(train_data_collated): + with torch.no_grad(): + out = gm(**frame, evaluation_mode=EvaluationMode.EVALUATION) + + image_rgb = to_numpy_image(out["images_render"][0]) + mask = to_numpy_image(out["masks_render"][0]) + expd = to_numpy_image(resize_image(frame.image_rgb)[0]) + mask_expected = to_numpy_image(resize_image(frame.fg_probability)[0]) + + images.append(image_rgb) + masks.append(mask) + expected.append(expd) + masks_expected.append(mask_expected) + + +# We draw a grid showing predicted image and expected image, followed by predicted mask and expected mask, from each viewpoint. +# This is a grid of four rows of images, wrapped in to several large rows, i.e.. +#
+# ``` +# ┌────────┬────────┐ ┌────────┐ +# │pred │pred │ │pred │ +# │image │image │ │image │ +# │1 │2 │ │n │ +# ├────────┼────────┤ ├────────┤ +# │expected│expected│ │expected│ +# │image │image │ ... │image │ +# │1 │2 │ │n │ +# ├────────┼────────┤ ├────────┤ +# │pred │pred │ │pred │ +# │mask │mask │ │mask │ +# │1 │2 │ │n │ +# ├────────┼────────┤ ├────────┤ +# │expected│expected│ │expected│ +# │mask │mask │ │mask │ +# │1 │2 │ │n │ +# ├────────┼────────┤ ├────────┤ +# │pred │pred │ │pred │ +# │image │image │ │image │ +# │n+1 │n+1 │ │2n │ +# ├────────┼────────┤ ├────────┤ +# │expected│expected│ │expected│ +# │image │image │ ... │image │ +# │n+1 │n+2 │ │2n │ +# ├────────┼────────┤ ├────────┤ +# │pred │pred │ │pred │ +# │mask │mask │ │mask │ +# │n+1 │n+2 │ │2n │ +# ├────────┼────────┤ ├────────┤ +# │expected│expected│ │expected│ +# │mask │mask │ │mask │ +# │n+1 │n+2 │ │2n │ +# └────────┴────────┘ └────────┘ +# ... +# ``` +#
+ +# In[ ]: + + +images_to_display = [images.copy(), expected.copy(), masks.copy(), masks_expected.copy()] +n_rows = 4 +n_images = len(images) +blank_image = images[0] * 0 +n_per_row = 1+(n_images-1)//n_rows +for _ in range(n_per_row*n_rows - n_images): + for group in images_to_display: + group.append(blank_image) + +images_to_display_listed = [[[i] for i in j] for j in images_to_display] +split = [] +for row in range(n_rows): + for group in images_to_display_listed: + split.append(group[row*n_per_row:(row+1)*n_per_row]) + +Image.fromarray(np.block(split)) + + +# In[ ]: + + +# Print the maximum channel intensity in the first image. +print(images[1].max()/255) + + +# In[ ]: + + +plt.ioff() +fig, ax = plt.subplots(figsize=(3,3)) + +ax.grid(None) +ims = [[ax.imshow(im, animated=True)] for im in images] +ani = animation.ArtistAnimation(fig, ims, interval=80, blit=True) +ani_html = ani.to_jshtml() + + +# In[ ]: + + +HTML(ani_html) + + +# In[ ]: + + +# If you want to see the output of the model with the volume forced to opaque white, run this and re-evaluate +# with torch.no_grad(): +# gm._implicit_functions[0]._fn.density.fill_(9.0) +# gm._implicit_functions[0]._fn.color.fill_(9.0) + diff --git a/files/render_colored_points.ipynb b/files/render_colored_points.ipynb index c28c35ce..23e9cd14 100644 --- a/files/render_colored_points.ipynb +++ b/files/render_colored_points.ipynb @@ -50,7 +50,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/render_colored_points.py b/files/render_colored_points.py index 4744271f..da277ff2 100644 --- a/files/render_colored_points.py +++ b/files/render_colored_points.py @@ -30,7 +30,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/render_densepose.ipynb b/files/render_densepose.ipynb index 910958c9..2e668776 100644 --- a/files/render_densepose.ipynb +++ b/files/render_densepose.ipynb @@ -57,7 +57,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/render_densepose.py b/files/render_densepose.py index c498f206..8e87ff14 100644 --- a/files/render_densepose.py +++ b/files/render_densepose.py @@ -34,7 +34,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/files/render_textured_meshes.ipynb b/files/render_textured_meshes.ipynb index 8de61ba1..333eb3b7 100644 --- a/files/render_textured_meshes.ipynb +++ b/files/render_textured_meshes.ipynb @@ -73,7 +73,7 @@ "except ModuleNotFoundError:\n", " need_pytorch3d=True\n", "if need_pytorch3d:\n", - " if torch.__version__.startswith(\"1.10.\") and sys.platform.startswith(\"linux\"):\n", + " if torch.__version__.startswith(\"1.12.\") 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", diff --git a/files/render_textured_meshes.py b/files/render_textured_meshes.py index 3550e778..7c35e849 100644 --- a/files/render_textured_meshes.py +++ b/files/render_textured_meshes.py @@ -32,7 +32,7 @@ try: except ModuleNotFoundError: need_pytorch3d=True if need_pytorch3d: - if torch.__version__.startswith("1.10.") and sys.platform.startswith("linux"): + if torch.__version__.startswith("1.12.") and sys.platform.startswith("linux"): # We try to install PyTorch3D via a released wheel. pyt_version_str=torch.__version__.split("+")[0].replace(".", "") version_str="".join([ diff --git a/tutorials/bundle_adjustment.html b/tutorials/bundle_adjustment.html index 7b8830e4..93e2671e 100644 --- a/tutorials/bundle_adjustment.html +++ b/tutorials/bundle_adjustment.html @@ -6,7 +6,7 @@ ga('create', 'UA-157376881-1', 'auto'); ga('send', 'pageview'); -