diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index 4e9f2ba85..0bb180492 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -121,6 +121,7 @@ class RFFEncoding(Encoding): if not scale > 0: raise ValueError("RFF encoding scale should be greater than zero") self.scale = scale + self.b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies)) def get_out_dim(self) -> int: return self.num_frequencies * 2 @@ -133,8 +134,7 @@ class RFFEncoding(Encoding): Returns: TensorType[..., "output_dim"]: Output values will be between -1 and 1 """ - b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies)) - scaled_inputs = in_tensor @ b_matrix # [..., "num_frequencies"] + scaled_inputs = in_tensor @ self.b_matrix # [..., "num_frequencies"] encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1) return encoded_inputs diff --git a/notebooks/visualize_encoders.ipynb b/notebooks/visualize_encoders.ipynb new file mode 100644 index 000000000..2fde5f60c --- /dev/null +++ b/notebooks/visualize_encoders.ipynb @@ -0,0 +1,287 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "import mediapy as media\n", + "from mattport.nerf.field_modules import encoding\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# NeRF Positional Encoding\n", + "First introduced in the original NeRF paper. This encoding assumes the inputs are between zero and one and can opperate on any dimensional input." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Input Values:\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Encoded Values:\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "num_frequencies = 4\n", + "min_freq_exp = 0\n", + "max_freq_exp = 6\n", + "include_input = True\n", + "resolution = 128\n", + "\n", + "encoder = encoding.NeRFEncoding(\n", + " in_dim=2,\n", + " num_frequencies=num_frequencies,\n", + " min_freq_exp=min_freq_exp,\n", + " max_freq_exp=max_freq_exp,\n", + " include_input=include_input,\n", + ")\n", + "\n", + "x_samples = torch.linspace(0, 1, resolution)\n", + "grid = torch.stack(torch.meshgrid([x_samples, x_samples], indexing=\"ij\"), dim=-1)\n", + "\n", + "encoded_values = encoder(grid)\n", + "\n", + "print(\"Input Values:\")\n", + "media.show_images(torch.moveaxis(grid, 2, 0), cmap=\"plasma\", border=True)\n", + "print(\"Encoded Values:\")\n", + "media.show_images(torch.moveaxis(encoded_values, 2, 0), cmap=\"plasma\", border=True)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Random Fourier Feature (RFF) Encoding\n", + "This encoding assumes the inputs are between zero and one and can opperate on any dimensional input." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Input Values:\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Encoded Values:\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "num_frequencies = 8\n", + "scale = 10\n", + "resolution = 128\n", + "\n", + "encoder = encoding.RFFEncoding(in_dim=2, num_frequencies=num_frequencies, scale=scale)\n", + "\n", + "x_samples = torch.linspace(0, 1, resolution)\n", + "grid = torch.stack(torch.meshgrid([x_samples, x_samples], indexing=\"ij\"), dim=-1)\n", + "\n", + "encoded_values = encoder(grid)\n", + "\n", + "print(\"Input Values:\")\n", + "media.show_images(torch.moveaxis(grid, 2, 0), cmap=\"plasma\", border=True)\n", + "print(\"Encoded Values:\")\n", + "media.show_images(torch.moveaxis(encoded_values, 2, 0), cmap=\"plasma\", border=True)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Hash Encoding\n", + "The hash incoding was originally introduced in Instant-NGP. The encoding is optimized during training. This is a visualization of the initialization." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Input Values:\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Encoded Values:\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "num_levels = 8\n", + "min_res = 2\n", + "max_res = 128\n", + "hash_table_size = 2**4 # Typically much larger tables are used\n", + "\n", + "resolution = 128\n", + "slice = 0\n", + "\n", + "# Fixing features_per_level to 3 for easy RGB visualization. Typical value is 2 in networks\n", + "features_per_level = 3\n", + "\n", + "encoder = encoding.HashEncoding(\n", + " num_levels=num_levels,\n", + " min_res=min_res,\n", + " max_res=max_res,\n", + " hash_table_size=hash_table_size,\n", + " features_per_level=features_per_level,\n", + " hash_init_scale=0.001,\n", + ")\n", + "\n", + "x_samples = torch.linspace(0, 1, resolution)\n", + "grid = torch.stack(torch.meshgrid([x_samples, x_samples, x_samples], indexing=\"ij\"), dim=-1)\n", + "\n", + "encoded_values = encoder(grid)\n", + "\n", + "grid_slice = grid[slice, ...]\n", + "encoded_values_slice = encoded_values[slice, ...]\n", + "\n", + "print(\"Input Values:\")\n", + "media.show_images(torch.moveaxis(grid_slice, 2, 0), cmap=\"plasma\", border=True)\n", + "\n", + "print(\"Encoded Values:\")\n", + "encoded_images = encoded_values_slice.view(resolution, resolution, num_levels, 3)\n", + "encoded_images = torch.moveaxis(encoded_images, 2, 0)\n", + "encoded_images -= torch.min(encoded_images)\n", + "encoded_images /= torch.max(encoded_images)\n", + "media.show_images(encoded_images.detach().numpy(), cmap=\"plasma\", border=True)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "interpreter": { + "hash": "024fe9dea615e4e34961116ff517c7e9f5910440e3d2c8eb4560299f11aaf2e7" + }, + "kernelspec": { + "display_name": "Python 3.9.11 ('mattport')", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.11" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +}