lora-nf 0.1.1__tar.gz

lora_nf-0.1.1/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: lora-nf
+Version: 0.1.1
+Summary: Low-Rank Adaptation of Neural Fields
+Home-page: https://github.com/dinhanhtruong/LoRA-NF
+Author: Anh Truong
+License: MIT License
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: imageio
+Requires-Dist: trimesh
+Requires-Dist: pysdf
+Requires-Dist: PyMCubes
+Requires-Dist: packaging
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: summary

lora_nf-0.1.1/README.md

@@ -0,0 +1,2 @@
+# LoRA-NF
+TODO

lora_nf-0.1.1/lora_nf/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.1"

lora_nf-0.1.1/lora_nf/common.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+
+# Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+# 
+# Redistribution and use in source and binary forms, with or without modification, are permitted
+# provided that the following conditions are met:
+#     * Redistributions of source code must retain the above copyright notice, this list of
+#       conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above copyright notice, this list of
+#       conditions and the following disclaimer in the documentation and/or other materials
+#       provided with the distribution.
+#     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+#       to endorse or promote products derived from this software without specific prior written
+#       permission.
+# 
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+# STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import imageio
+import numpy as np
+import os
+import struct
+
+ROOT_DIR = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+
+def mse2psnr(x):
+	return -10.*np.log(x)/np.log(10.)
+
+def write_image_imageio(img_file, img, quality):
+	img = (np.clip(img, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)   # 
+	kwargs = {}
+	if os.path.splitext(img_file)[1].lower() in [".jpg", ".jpeg"]:
+		if img.ndim >= 3 and img.shape[2] > 3:
+			img = img[:,:,:3]
+		if img.ndim >= 3 and img.shape[2] == 1:  # grayscale 
+			img = np.squeeze(img)
+		kwargs["quality"] = quality
+		kwargs["subsampling"] = 0
+	imageio.imwrite(img_file, img, **kwargs)
+
+def read_image_imageio(img_file):
+	img = imageio.imread(img_file)
+	img = np.asarray(img).astype(np.float32)
+	if len(img.shape) == 2:
+		img = img[:,:,np.newaxis]
+	return img / 255.0
+
+def srgb_to_linear(img):
+	limit = 0.04045
+	return np.where(img > limit, np.power((img + 0.055) / 1.055, 2.4), img / 12.92)
+
+def linear_to_srgb(img):
+	limit = 0.0031308
+	return np.where(img > limit, 1.055 * (img ** (1.0 / 2.4)) - 0.055, 12.92 * img)
+
+def read_image(file):
+	if os.path.splitext(file)[1] == ".bin":
+		with open(file, "rb") as f:
+			bytes = f.read()
+			h, w = struct.unpack("ii", bytes[:8])
+			img = np.frombuffer(bytes, dtype=np.float16, count=h*w*4, offset=8).astype(np.float32).reshape([h, w, 4])
+	else:
+		img = read_image_imageio(file)
+		if img.shape[2] == 4:
+			img[...,0:3] = srgb_to_linear(img[...,0:3])
+			# Premultiply alpha
+			img[...,0:3] *= img[...,3:4]
+		else:
+			img = srgb_to_linear(img)
+	return img
+
+def write_image(file, img, quality=100):
+	if os.path.splitext(file)[1] == ".bin":
+		if img.shape[2] < 4:
+			img = np.dstack((img, np.ones([img.shape[0], img.shape[1], 4 - img.shape[2]])))
+		with open(file, "wb") as f:
+			f.write(struct.pack("ii", img.shape[0], img.shape[1]))
+			f.write(img.astype(np.float16).tobytes())
+	else:
+		if img.shape[2] == 4:
+			img = np.copy(img)
+			# Unmultiply alpha
+			img[...,0:3] = np.divide(img[...,0:3], img[...,3:4], out=np.zeros_like(img[...,0:3]), where=img[...,3:4] != 0)
+			img[...,0:3] = linear_to_srgb(img[...,0:3])
+		else:
+			img = linear_to_srgb(img)
+		write_image_imageio(file, img, quality)
+
+def trim(error, skip=0.000001):
+	error = np.sort(error.flatten())
+	size = error.size
+	skip = int(skip * size)
+	return error[skip:size-skip].mean()
+
+def luminance(a):
+	a = np.maximum(0, a)**0.4545454545
+	return 0.2126 * a[:,:,0] + 0.7152 * a[:,:,1] + 0.0722 * a[:,:,2]
+
+def L1(img, ref):
+	return np.abs(img - ref)
+
+def APE(img, ref):
+	return L1(img, ref) / (1e-2 + ref)
+
+def SAPE(img, ref):
+	return L1(img, ref) / (1e-2 + (ref + img) / 2.)
+
+def L2(img, ref):
+	return (img - ref)**2
+
+def RSE(img, ref):
+	return L2(img, ref) / (1e-2 + ref**2)
+
+def rgb_mean(img):
+	return np.mean(img, axis=2)
+
+def compute_error_img(metric, img, ref):
+	img[np.logical_not(np.isfinite(img))] = 0
+	img = np.maximum(img, 0.)
+	if metric == "MAE":
+		return L1(img, ref)
+	elif metric == "MAPE":
+		return APE(img, ref)
+	elif metric == "SMAPE":
+		return SAPE(img, ref)
+	elif metric == "MSE":
+		return L2(img, ref)
+	elif metric == "MScE":
+		return L2(np.clip(img, 0.0, 1.0), np.clip(ref, 0.0, 1.0))
+	elif metric == "MRSE":
+		return RSE(img, ref)
+	elif metric == "MtRSE":
+		return trim(RSE(img, ref))
+	elif metric == "MRScE":
+		return RSE(np.clip(img, 0, 100), np.clip(ref, 0, 100))
+
+	raise ValueError(f"Unknown metric: {metric}.")
+
+def compute_error(metric, img, ref):
+	metric_map = compute_error_img(metric, img, ref)
+	metric_map[np.logical_not(np.isfinite(metric_map))] = 0
+	if len(metric_map.shape) == 3:
+		metric_map = np.mean(metric_map, axis=2)
+	mean = np.mean(metric_map)
+	return mean

lora_nf-0.1.1/lora_nf/data_samplers.py

@@ -0,0 +1,169 @@
+import os
+import numpy as np
+import torch
+from .common import read_image, write_image
+import trimesh
+from .util import save_mesh
+import pysdf
+
+class DataSampler:
+    def sample_batch(self, n_samples: int, device: torch.device):
+        """
+        samples input positions and corresponding field values
+        
+        Args
+            - n_samples: number of samples
+            - device: torch device (e.g. "cpu", "cuda:0") on which the sampled data will live
+
+        Returns:
+            - inputs: sample positions. torch tensor of shape [n_samples, data_input_dim], e.g., [n_samples, 2] for image xy positions
+            - targets: values of the field sampled at inputs. torch tensor of shape [n_samples, data_output_dim], e.g., [n_samples, 3] for image RGB values
+        """
+        raise NotImplementedError
+    def save_model_output(self, model: torch.nn.Module, save_path: str):
+        """
+        saves a reconstruction of the given neural field at the given path (without extension). assumes that the directory exists
+
+        Args: 
+            - model: a neural field nn.Module()
+        
+        Returns nothing
+        """
+        raise NotImplementedError
+    
+class Image(DataSampler, torch.nn.Module):
+    def __init__(self, filename: str, device: torch.device):
+        """
+        Args
+            - filename: path to image
+            - device: torch.device
+        """
+        super().__init__()
+        self.data = read_image(filename)
+        # remove alpha channel
+        if self.data.ndim > 2 and self.data.shape[2] == 4:
+            self.data = self.data[:,:,:3] # keep RGB
+        self.orig_data_npy = self.data
+        self.shape = self.data.shape
+        self.data = torch.from_numpy(self.data).float().to(device)
+        self.device = device
+
+        # for model output reconsturction
+        resolution = self.data.shape[0:2]
+        # n_pixels = resolution[0] * resolution[1]
+        half_dx =  0.5 / resolution[0] # half pixel size
+        half_dy =  0.5 / resolution[1]
+        xs = torch.linspace(half_dx, 1-half_dx, resolution[0])
+        ys = torch.linspace(half_dy, 1-half_dy, resolution[1])
+        xv, yv = torch.meshgrid([xs, ys], indexing="ij")
+        self.img_shape = resolution + torch.Size([self.data.shape[2]])
+        self.xy = torch.stack((yv.flatten(), xv.flatten())).t().to(device)
+
+    def forward(self, xs, interpolate=True):
+        with torch.no_grad():
+            # Bilinearly filtered lookup from the image. Not super fast,
+            # but less than ~20% of the overall runtime of this example.
+            shape = self.shape
+
+            xs = xs * torch.tensor([shape[1], shape[0]], device=xs.device).float()
+            indices = xs.long()
+
+            x0 = indices[:, 0].clamp(min=0, max=shape[1]-1)
+            y0 = indices[:, 1].clamp(min=0, max=shape[0]-1)
+            if interpolate:
+                lerp_weights = xs - indices.float()
+                x1 = (x0 + 1).clamp(max=shape[1]-1)
+                y1 = (y0 + 1).clamp(max=shape[0]-1)
+
+                return (
+                    self.data[y0, x0] * (1.0 - lerp_weights[:,0:1]) * (1.0 - lerp_weights[:,1:2]) +
+                    self.data[y0, x1] * lerp_weights[:,0:1] * (1.0 - lerp_weights[:,1:2]) +
+                    self.data[y1, x0] * (1.0 - lerp_weights[:,0:1]) * lerp_weights[:,1:2] +
+                    self.data[y1, x1] * lerp_weights[:,0:1] * lerp_weights[:,1:2]
+                )
+            
+            return (self.data[y0, x0]) # no interpolation
+    
+    def sample_batch(self, n_samples, device):
+        assert device == self.device
+        input_xy = torch.rand([n_samples, 2], dtype=torch.float, device=device) 
+        image_rgb = self.forward(input_xy)
+        return input_xy, image_rgb
+
+    def save_model_output(self, model, save_path):
+        write_image(f"{save_path}.png", model(self.xy).reshape(self.img_shape).clamp(0.0, 1.0).detach().cpu().numpy())
+
+
+class SDF(DataSampler):
+    def __init__(self, path: str, device: torch.device, num_samples=2**18, clip_sdf=None, transformation_save_dir=""):
+        super().__init__()
+        self.path = path
+        self.device = self.device
+
+        # load obj 
+        self.mesh = trimesh.load(path, force='mesh')
+
+        # normalize to [-1, 1] (different from instant-sdf where is [0, 1]) via scaling and translation
+        vs = self.mesh.vertices
+        vmin = vs.min(0)
+        vmax = vs.max(0)
+        v_center = (vmin + vmax) / 2
+        v_scale = 2 / np.sqrt(np.sum((vmax - vmin) ** 2)) * 0.95
+
+        # TODO: save normalizing transformation of base model and apply same transformation to deformed model 
+        transformation_path = f"{transformation_save_dir}/data_normalization_transformation.npz"
+        if not os.path.exists(transformation_path):
+            print("saving base mesh normalization transformation")
+            np.savez(transformation_path, v_scale=v_scale, v_center=v_center)
+        else:
+            print("######################")
+            print("####### TEMP DISABLED CONSISTENT SDF NORMALIZATION")
+            print("######################")
+            print("######################")
+            print("######################")
+            print("######################")
+        #     print("using existing normalization transformation")
+        #     loaded_transformation = np.load(transformation_path)
+        #     v_scale = loaded_transformation["v_scale"]
+        #     v_center = loaded_transformation["v_center"]
+
+        print("scale: ", v_scale)
+        print("center: ", v_center)
+        # apply transformation to verts
+        vs = (vs - v_center[None, :]) * v_scale
+        self.mesh.vertices = vs
+
+        print(f"[INFO] mesh verts & faces: {self.mesh.vertices.shape} & {self.mesh.faces.shape}")
+
+        if not self.mesh.is_watertight:
+            print(f"[WARN] mesh is not watertight! SDF maybe incorrect.")
+        #trimesh.Scene([self.mesh]).show()
+
+        self.sdf_fn = pysdf.SDF(self.mesh.vertices, self.mesh.faces)
+        
+        self.num_samples = num_samples
+        assert self.num_samples % 8 == 0, "num_samples must be divisible by 8."
+        self.clip_sdf = clip_sdf
+
+    def sample_batch(self, n_samples, device):
+        assert device == self.device
+        # online sampling
+        sdfs = np.zeros((n_samples, 1))
+        # surface query points (7/8 points for surface and near-surface)
+        points_surface = self.mesh.sample(n_samples * 7 // 8)
+        
+        # near-surface points
+        points_surface[n_samples // 2:] += 0.01 * np.random.randn(n_samples * 3 // 8, 3)
+        # random uniform points (1/8 of points)
+        points_uniform = np.random.rand(n_samples // 8, 3) * 2 - 1
+        points = np.concatenate([points_surface, points_uniform], axis=0).astype(np.float32)
+        sdfs[n_samples // 2:] = -self.sdf_fn(points[n_samples // 2:])[:,None].astype(np.float32)
+ 
+        # clip sdf
+        if self.clip_sdf is not None:
+            sdfs = sdfs.clip(-self.clip_sdf, self.clip_sdf)
+
+        return torch.tensor(points, device=device), torch.tensor(sdfs, device=device)
+
+    def save_model_output(self, model: torch.nn.Module, save_path: str):
+        save_mesh(save_path, model, self.device, resolution=256)

lora_nf-0.1.1/lora_nf/torch_modules.py

@@ -0,0 +1,122 @@
+import math
+import torch
+import torch.nn as nn
+
+class CustomFrequencyEncoding(nn.Module):
+    def __init__(self):
+        super(CustomFrequencyEncoding, self).__init__()
+    def get_encoding_output_dim(self, input_dim):
+        return self.forward(torch.zeros((1, input_dim))).shape[-1]
+
+    def forward( # from https://github.com/krrish94/nerf-pytorch/blob/master/nerf/nerf_helpers.py
+        self, tensor, num_encoding_functions=10, include_input=True, log_sampling=True
+    ) -> torch.Tensor:
+        r"""Apply positional encoding to the input.
+
+        Args:
+            tensor (torch.Tensor): Input tensor to be positionally encoded.
+            encoding_size (optional, int): Number of encoding functions used to compute
+                a positional encoding (default: 6).
+            include_input (optional, bool): Whether or not to include the input in the
+                positional encoding (default: True).
+
+        Returns:
+        (torch.Tensor): Positional encoding of the input tensor.
+        """
+        encoding = [tensor] if include_input else []
+        frequency_bands = None
+        if log_sampling:
+            frequency_bands = 2.0 ** torch.linspace(
+                0.0,
+                num_encoding_functions - 1,
+                num_encoding_functions,
+                dtype=tensor.dtype,
+                device=tensor.device,
+            )
+        else:
+            frequency_bands = torch.linspace(
+                2.0 ** 0.0,
+                2.0 ** (num_encoding_functions - 1),
+                num_encoding_functions,
+                dtype=tensor.dtype,
+                device=tensor.device,
+            )
+
+        for freq in frequency_bands:
+            for func in [torch.sin, torch.cos]:
+                encoding.append(func(tensor * freq))
+
+        # Special case, for no positional encoding
+        if len(encoding) == 1:
+            return encoding[0]
+        else:
+            return torch.cat(encoding, dim=-1)
+
+class LoRA_MLP(nn.Module):
+    def __init__(self, base_mlp, rank):
+        '''
+        LoRA-augmented MLP 
+
+        prev_frame_lora_paths: list of chronologically ordered filepaths to prev frames' LoraMLP weights 
+        '''
+        super(LoRA_MLP, self).__init__()
+
+        # for each parent MLP linear layer, store a low-rank adaptor (LoRALinear layer).
+        # mimic parent MLP but replace linear with LoRALinear
+        self.sequential = nn.Sequential() # just LoRALinears
+
+        for base_layer in base_mlp:
+            if isinstance(base_layer, nn.Linear):
+                # add lora layer
+                self.sequential.append(LoRALinear(base_layer, r=rank))
+            else:
+                # keep parent layer (activation or positional encoding)
+                self.sequential.append(base_layer)
+  
+    def as_sequential(self):
+        return self.sequential
+    def get_lora_weights(self):
+        """
+        Returns a list of tensors of weights corresponding to the current LoRAs (in the same order as the base model's)
+        """
+        lora_weights = []
+        for lora_mlp_layer in self.sequential:
+            if isinstance(lora_mlp_layer, LoRALinear):
+                lora_weights.append((lora_mlp_layer.A @ lora_mlp_layer.B).T) # [in_features, out_features]
+        return lora_weights
+
+    def forward(self, x):
+        return self.sequential(x)
+    
+
+class LoRALinear(nn.Module):
+    def __init__(self, base_linear, r=16, alpha=1, A=None):
+        '''
+        
+        '''
+        super(LoRALinear, self).__init__()
+        assert isinstance(base_linear, nn.Linear) and r > 1
+        self.base_linear = base_linear
+        # freeze base; make sure only lora weights are trainable
+        self.base_linear.weight.requires_grad_(False)
+        if self.base_linear.bias is not None:
+            self.base_linear.bias.requires_grad_(False)
+
+        rank = min(r, base_linear.in_features, base_linear.out_features)
+        if A is not None:
+        # use specified A matrix
+            assert A.shape == torch.Size((base_linear.in_features, rank))
+            self.A = A
+        else:
+            self.A = nn.Parameter(torch.empty((base_linear.in_features, rank)))
+            nn.init.normal_(self.A, std=1/math.sqrt(base_linear.in_features)) #following pg4 of https://arxiv.org/pdf/2406.08447 
+        self.B = nn.Parameter(torch.zeros(rank, base_linear.out_features))
+        # lora scale factor
+        self.scaling = alpha/r
+
+    def forward(self, x):
+        return self.base_linear(x) + torch.linalg.multi_dot((x, self.A, self.B))*self.scaling
+    
+
+def extract_linear_layers(module):
+    return [m for m in module.modules() if isinstance(m, nn.Linear)]

lora_nf-0.1.1/lora_nf/train_lora.py

@@ -0,0 +1,198 @@
+import os
+import random
+import numpy as np
+import time
+import torch
+import copy
+import torch.nn as nn
+from .torch_modules import LoRA_MLP, extract_linear_layers
+from .util import get_device
+from .data_samplers import DataSampler
+from typing import Callable
+
+torch.manual_seed(0)
+np.random.seed(0)
+random.seed(0)
+
+
+def train_lora_regression(
+        base_nf: nn.Sequential, 
+        target_sampler: DataSampler, 
+        loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], 
+        lora_rank: int, 
+        learning_rate=5e-3, 
+        batch_size=2**18, 
+        max_n_steps=30000, 
+        lr_scheduler_warmup_steps=7000, 
+        log_interval=100,
+        convergence_patience=15,
+        save_dir="",
+    ):
+    """
+    Trains LoRAs for every linear layer of the given base neural field. Returns a new neural field with a LoRA applied and (optionally) the weights of each LoRA
+    
+    Args:
+        - base_nf: MLP implemented as a nn.Sequential (containing nn.Linear layers, activation functions, optional input positional encoding)
+        - target_sampler: a DataSampler that must be compatible with base_nf (i.e. input/target dimensions must match)
+        - loss_fn: callable function (output,target) |-> loss scalar
+        - lora_rank: Desired maximum rank of each LoRA
+        - learning_rate: for ADAM optimizer
+        - batch_size: num samples per step
+        - max_n_steps: max number of training steps
+        - lr_scheduler_warmup_steps: number of warmup steps for the learning rate scheduler. No warmup if 0.
+        - convergence_patience: number of log_intervals of no improvement after which training is terminated early
+        - save_dir: if provided, then weights and output reconstructions will be saved there
+    
+    Returns:
+        - lora_weights: LoRA weight tensors, one per linear layer of the base model. Returns the iterate with the lowest loss.
+        - lora_nf: nn.Sequential neural field with LoRA applied to every linear layer. Returns the iterate with the lowest loss.
+    """
+    # make sure that base_nf and target_sampler are compatible
+    device = get_device(base_nf)
+    linear_layers = extract_linear_layers(base_nf)
+    base_output_dim = linear_layers[-1].out_features
+    _, dummy_targets = target_sampler.sample_batch(batch_size, device)
+    assert dummy_targets.shape[-1] == base_output_dim, "base_nf and target_sampler are incompatible: output size mismatch"
+    assert isinstance(base_nf, nn.Sequential)
+    
+    # set up lora neural field
+    lora_nf = LoRA_MLP(base_nf, lora_rank)
+    lora_nf.to(device)
+    print(f"training LoRA on {device}")
+    trainable_params = sum(p.numel() for p in lora_nf.parameters() if p.requires_grad)
+    print(f"# trainable LoRA parameters: {trainable_params}")
+
+    optimizer = torch.optim.Adam(lora_nf.parameters(), lr=learning_rate) # only fine-tune mlp
+    # use learning rate scheduler if specified
+    scheduler = None        
+    if lr_scheduler_warmup_steps > 0:
+        scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, total_iters=lr_scheduler_warmup_steps)
+    # for logging
+    best_loss = float('inf')
+    best_lora_nf = copy.deepcopy(lora_nf)
+    prev_time = time.perf_counter()
+    periods_no_improve = 0 # for early stopping
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+
+    for i in range(max_n_steps):
+        batch_inputs, batch_targets = target_sampler.sample_batch(batch_size, device) # [B,in], [B,out]
+        output = lora_nf(batch_inputs.to(device))
+        loss = loss_fn(output, batch_targets.to(device))
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        if scheduler:
+            scheduler.step()
+
+        if i % log_interval == 0: # reconstruct output
+            curr_loss = loss.item()
+            elapsed_time = time.perf_counter() - prev_time
+            print(f"Step#{i}: loss={curr_loss:.7f} time={int(elapsed_time)}[s]")
+            if save_dir:
+                with torch.no_grad():
+                    target_sampler.save_model_output(lora_nf, save_path=f"{save_dir}/lora_nf_step_{i:05d}")
+
+            if curr_loss < best_loss:
+                print(f"\tdecreased {best_loss:.6f}-->{curr_loss:.6f}")
+                best_loss = curr_loss
+                periods_no_improve = 0 # reset
+                best_lora_nf = copy.deepcopy(lora_nf)
+            else:
+                # early stopping if no improvement for several epochs
+                periods_no_improve += 1
+                if periods_no_improve >= convergence_patience:
+                    print(f"Early stopping at step {i} with loss {curr_loss}")
+                    break
+            prev_time = time.perf_counter()
+            if i > 0 and log_interval < 1000: 
+                log_interval *= 10
+    if save_dir:
+        torch.save({
+            'step': i,
+            'model_state_dict': best_lora_nf.state_dict(),
+            'loss': curr_loss,
+        }, f"{save_dir}/lora_nf_best.pt")
+        target_sampler.save_model_output(best_lora_nf, save_path=f"{save_dir}/lora_nf_best")
+        
+    return best_lora_nf.get_lora_weights(), best_lora_nf.as_sequential()
+
+def train_base_model(
+        base_nf: nn.Sequential, 
+        data_sampler: DataSampler, 
+        loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], 
+        learning_rate=1e-4, 
+        batch_size=2**18, 
+        max_n_steps=100000, 
+        lr_scheduler_warmup_steps=0, 
+        log_interval=100,
+        convergence_patience=15,
+        save_dir="",
+    ):
+    """
+    Trains the given base neural field to regress samples from data_sampler.
+
+    Returns a copy of base_nf at the best training iteration
+    """
+    device = get_device(base_nf)
+    trainable_params = sum(p.numel() for p in base_nf.parameters() if p.requires_grad)
+    print(f"training base model on {device}")
+    print(f"# trainable base model parameters: {trainable_params}")
+    
+    optimizer = torch.optim.Adam(base_nf.parameters(), lr=learning_rate) # only fine-tune mlp
+    # use learning rate scheduler if specified
+    scheduler = None        
+    if lr_scheduler_warmup_steps > 0:
+        scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, total_iters=lr_scheduler_warmup_steps)
+    best_loss = float('inf')
+    best_model = copy.deepcopy(base_nf)
+    prev_time = time.perf_counter()
+    periods_no_improve = 0 # for early stopping
+
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+
+    for i in range(max_n_steps):
+        batch_inputs, batch_targets = data_sampler.sample_batch(batch_size, device) # [B,in], [B,out]
+        output = base_nf(batch_inputs.to(device))
+        loss = loss_fn(output, batch_targets.to(device))
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        if scheduler:
+            scheduler.step()
+
+        if i % log_interval == 0: # reconstruct output
+            curr_loss = loss.item()
+            elapsed_time = time.perf_counter() - prev_time
+            print(f"Step#{i}: loss={curr_loss:.7f} time={int(elapsed_time)}[s]")
+            if save_dir:
+                with torch.no_grad():
+                    data_sampler.save_model_output(base_nf, save_path=f"{save_dir}/base_nf_step_{i:05d}")
+
+            if curr_loss < best_loss:
+                print(f"\tdecreased {best_loss:.6f}-->{curr_loss:.6f}")
+                best_loss = curr_loss
+                periods_no_improve = 0 # reset
+                best_model = copy.deepcopy(base_nf)
+            else:
+                # early stopping if no improvement for several epochs
+                periods_no_improve += 1
+                if periods_no_improve >= convergence_patience:
+                    print(f"Early stopping at step {i} with loss {curr_loss}")
+                    break
+            prev_time = time.perf_counter()
+            if i > 0 and log_interval < 1000: 
+                log_interval *= 10
+    if save_dir:
+        data_sampler.save_model_output(best_model, save_path=f"{save_dir}/base_nf_best")
+        torch.save({
+            'step': i,
+            'model_state_dict': best_model.state_dict(),
+            'loss': curr_loss,
+        }, f"{save_dir}/base_nf_best.pt")
+
+    return best_model
+

lora_nf-0.1.1/lora_nf/util.py

@@ -0,0 +1,269 @@
+import os
+import mcubes
+import numpy as np
+import packaging
+import torch
+import trimesh
+
+def get_device(module):
+    try:
+        return next(module.parameters()).device
+    except StopIteration:
+        return next(module.buffers()).device
+
+def check_shape_equality(*images): # borrowed from scikit-image
+    """Check that all images have the same shape"""
+    image0 = images[0]
+    if not all(image0.shape == image.shape for image in images[1:]):
+        raise ValueError('Input images must have the same dimensions.')
+    return
+
+
+#################
+## borrowed from scikit-image for psnr metric
+dtype_range = {
+    bool: (False, True),
+    np.bool_: (False, True),
+    float: (-1, 1),
+    np.float16: (-1, 1),
+    np.float32: (-1, 1),
+    np.float64: (-1, 1),
+}
+new_float_type = {
+    # preserved types
+    np.float32().dtype.char: np.float32,
+    np.float64().dtype.char: np.float64,
+    np.complex64().dtype.char: np.complex64,
+    np.complex128().dtype.char: np.complex128,
+    # altered types
+    np.float16().dtype.char: np.float32,
+    'g': np.float64,  # np.float128 ; doesn't exist on windows
+    'G': np.complex128,  # np.complex256 ; doesn't exist on windows
+}
+def _supported_float_type(input_dtype, allow_complex=False):
+    """Return an appropriate floating-point dtype for a given dtype.
+
+    float32, float64, complex64, complex128 are preserved.
+    float16 is promoted to float32.
+    complex256 is demoted to complex128.
+    Other types are cast to float64.
+
+    Parameters
+    ----------
+    input_dtype : np.dtype or tuple of np.dtype
+        The input dtype. If a tuple of multiple dtypes is provided, each
+        dtype is first converted to a supported floating point type and the
+        final dtype is then determined by applying `np.result_type` on the
+        sequence of supported floating point types.
+    allow_complex : bool, optional
+        If False, raise a ValueError on complex-valued inputs.
+
+    Returns
+    -------
+    float_type : dtype
+        Floating-point dtype for the image.
+    """
+    if isinstance(input_dtype, tuple):
+        return np.result_type(*(_supported_float_type(d) for d in input_dtype))
+    input_dtype = np.dtype(input_dtype)
+    if not allow_complex and input_dtype.kind == 'c':
+        raise ValueError("complex valued input is not supported")
+    return new_float_type.get(input_dtype.char, np.float64)
+def _as_floats(image0, image1):
+    """
+    Promote im1, im2 to nearest appropriate floating point precision.
+    """
+    float_type = _supported_float_type((image0.dtype, image1.dtype))
+    image0 = np.asarray(image0, dtype=float_type)
+    image1 = np.asarray(image1, dtype=float_type)
+    return image0, image1
+def mean_squared_error(image0, image1):
+    """
+    Compute the mean-squared error between two images.
+
+    Parameters
+    ----------
+    image0, image1 : ndarray
+        Images.  Any dimensionality, must have same shape.
+
+    Returns
+    -------
+    mse : float
+        The mean-squared error (MSE) metric.
+
+    Notes
+    -----
+    .. versionchanged:: 0.16
+        This function was renamed from ``skimage.measure.compare_mse`` to
+        ``skimage.metrics.mean_squared_error``.
+
+    """
+    check_shape_equality(image0, image1)
+    image0, image1 = _as_floats(image0, image1)
+    return np.mean((image0 - image1) ** 2, dtype=np.float64)
+
+def peak_signal_noise_ratio(image_true, image_test, *, data_range=None):
+    """
+    Compute the peak signal to noise ratio (PSNR) for an image.
+
+    Parameters
+    ----------
+    image_true : ndarray
+        Ground-truth image, same shape as im_test.
+    image_test : ndarray
+        Test image.
+    data_range : int, optional
+        The data range of the input image (distance between minimum and
+        maximum possible values).  By default, this is estimated from the image
+        data-type.
+
+    Returns
+    -------
+    psnr : float
+        The PSNR metric.
+
+    Notes
+    -----
+    .. versionchanged:: 0.16
+        This function was renamed from ``skimage.measure.compare_psnr`` to
+        ``skimage.metrics.peak_signal_noise_ratio``.
+
+    References
+    ----------
+    .. [1] https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
+
+    """
+    check_shape_equality(image_true, image_test)
+
+    if data_range is None:
+        if image_true.dtype != image_test.dtype:
+            print(
+                "Inputs have mismatched dtype.  Setting data_range based on "
+                "image_true."
+            )
+        dmin, dmax = dtype_range[image_true.dtype.type]
+        true_min, true_max = np.min(image_true), np.max(image_true)
+        if true_max > dmax or true_min < dmin:
+            raise ValueError(
+                "image_true has intensity values outside the range expected "
+                "for its data type. Please manually specify the data_range."
+            )
+        if true_min >= 0:
+            # most common case (255 for uint8, 1 for float)
+            data_range = dmax
+        else:
+            data_range = dmax - dmin
+
+    image_true, image_test = _as_floats(image_true, image_test)
+
+    err = mean_squared_error(image_true, image_test)
+    data_range = float(data_range)  # prevent overflow for small integer types
+    return 10 * np.log10((data_range**2) / err)
+#################
+#################
+
+
+def measure_gpu_memory(device):
+    if torch.cuda.is_available():
+        allocated_memory = torch.cuda.memory_allocated(device)
+        max_allocated_memory = torch.cuda.max_memory_allocated(device)
+        print(f"Allocated memory: {allocated_memory / 1024**2:.2f} MB")
+        # print(f"Max allocated memory: {max_allocated_memory / 1024**2:.2f} MB")
+
+
+def get_model_checkpoint_size_mb(checkpoint_path, return_param_count=False):
+    '''
+    checkpoint_path: full path containing file extension.
+    '''
+    checkpoint = torch.load(checkpoint_path, weights_only=True)["model_state_dict"]
+    size_model = 0
+    # for k,v in checkpoint.items():
+    #     print(k, v.shape)
+    param_count = 0
+    for param in checkpoint.values():
+        # if param.is_floating_point():
+        size_model += param.nelement() * param.element_size()
+        param_count += param.nelement()
+        # else:
+            # size_model += param.numel() * torch.iinfo(param.dtype).bits
+    print(f"\tmodel size: {(size_model / 1e6):.3f} MB")
+    print(f"\tnum params: {param_count}")
+    if return_param_count:
+        return size_model / 1e6, param_count
+    return size_model / 1e6
+
+
+def mean_relative_l2(pred, target, eps=0.01):
+    loss = (pred - target.to(pred.dtype))**2 / (pred.detach()**2 + eps)
+    return loss.mean()
+
+def mape_loss(pred, target, reduction='mean'):
+    # pred, target: [B, 1], torch tenspr
+    difference = (pred - target).abs()
+    scale = 1 / (target.abs() + 1e-2)
+    loss = difference * scale
+
+    if reduction == 'mean':
+        loss = loss.mean()
+    
+    return loss
+
+
+### marching cubes helpers from torch-ngp ###
+def custom_meshgrid(*args):
+    # ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid
+    if packaging.version.parse(torch.__version__) < packaging.version.parse('1.10'):
+        return torch.meshgrid(*args)
+    else:
+        return torch.meshgrid(*args, indexing='ij')
+
+def extract_fields(bound_min, bound_max, resolution, query_func, device=torch.device("cpu")):
+    N = 64
+    X = torch.linspace(bound_min[0], bound_max[0], resolution).to(device).split(N)
+    Y = torch.linspace(bound_min[1], bound_max[1], resolution).to(device).split(N)
+    Z = torch.linspace(bound_min[2], bound_max[2], resolution).to(device).split(N)
+
+    u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+    with torch.no_grad():
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = custom_meshgrid(xs, ys, zs)
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3]
+                    val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [N, 1] --> [x, y, z]
+                    u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
+    return u
+
+def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
+    u = extract_fields(bound_min, bound_max, resolution, query_func)
+    vertices, triangles = mcubes.marching_cubes(u, threshold)
+
+    b_max_np = bound_max.detach().cpu().numpy()
+    b_min_np = bound_min.detach().cpu().numpy()
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles
+
+
+
+def save_mesh(save_path, sdf_model, device, resolution=256):
+    print(f"==> Saving mesh to {save_path}")
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+    def query_func(pts):
+        pts = pts.to(device)
+        with torch.no_grad():
+            # with torch.cuda.amp.autocast(enabled=False):
+            sdfs = sdf_model(pts)
+        return sdfs
+
+    bounds_min = torch.FloatTensor([-1, -1, -1])
+    bounds_max = torch.FloatTensor([1, 1, 1])
+
+    vertices, triangles = extract_geometry(bounds_min, bounds_max, resolution=resolution, threshold=0, query_func=query_func)
+
+    mesh = trimesh.Trimesh(vertices, triangles, process=False) # important, process=True leads to seg fault...
+    mesh.export(save_path)
+
+    print(f"==> Finished saving mesh.")

lora_nf-0.1.1/lora_nf.egg-info/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: lora-nf
+Version: 0.1.1
+Summary: Low-Rank Adaptation of Neural Fields
+Home-page: https://github.com/dinhanhtruong/LoRA-NF
+Author: Anh Truong
+License: MIT License
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: imageio
+Requires-Dist: trimesh
+Requires-Dist: pysdf
+Requires-Dist: PyMCubes
+Requires-Dist: packaging
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: summary

lora_nf-0.1.1/lora_nf.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+README.md
+setup.py
+lora_nf/__init__.py
+lora_nf/common.py
+lora_nf/data_samplers.py
+lora_nf/torch_modules.py
+lora_nf/train_lora.py
+lora_nf/util.py
+lora_nf.egg-info/PKG-INFO
+lora_nf.egg-info/SOURCES.txt
+lora_nf.egg-info/dependency_links.txt
+lora_nf.egg-info/requires.txt
+lora_nf.egg-info/top_level.txt

lora_nf-0.1.1/lora_nf.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

lora_nf-0.1.1/lora_nf.egg-info/requires.txt

@@ -0,0 +1,7 @@
+torch
+numpy
+imageio
+trimesh
+pysdf
+PyMCubes
+packaging

lora_nf-0.1.1/lora_nf.egg-info/top_level.txt

@@ -0,0 +1 @@
+lora_nf

lora_nf-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

lora_nf-0.1.1/setup.py

@@ -0,0 +1,21 @@
+from setuptools import setup
+
+setup(
+    name='lora-nf',
+    version='0.1.1',    
+    description='Low-Rank Adaptation of Neural Fields',
+    url='https://github.com/dinhanhtruong/LoRA-NF',
+    author='Anh Truong',
+    license='MIT License',
+    packages=['lora_nf'],
+    install_requires=['torch',
+                      'numpy',                     
+                      'imageio',                     
+                      'trimesh',                     
+                      'pysdf',                     
+                      'PyMCubes',                     
+                      'packaging'],
+    classifiers=[
+        'Programming Language :: Python :: 3.12',
+    ],
+)