import torch
import torch.nn as nn
import numpy as np
from time import time
from torch.utils.data import Dataset, DataLoader

class NerfIntegrator:
    def __call__(self, rgb_values, sigma_values, interval_sizes):
        # the format of rgb_values is [...,number_of_samples,3]
        # the format if sigma_values is [...,number_of_samples,1]
        # where ... represents any amount of dimensions

        # the integral of sigma
        scaled_sigma = sigma_values * interval_sizes
        T_t = torch.exp(-torch.cumsum(scaled_sigma, dim=-1))
        colors = ((T_t * (1 - torch.exp(-scaled_sigma))).unsqueeze(-1) * rgb_values).sum(dim=-2)
        return colors

class ClassicalMLP(nn.Module):
    def __init__(self, number_of_inputs, number_of_outputs, layers_size, number_of_layers, activition_function=nn.ReLU):
        super().__init__()
        layers_channels = [number_of_inputs] + [layers_size for _ in
                                                range(number_of_layers)] + [number_of_outputs]
        layers = []
        for i, (input, output) in enumerate(zip(layers_channels[:-1], layers_channels[1:])):
            layers += [nn.Linear(input, output)]
            if i != len(layers_channels) - 2:
                layers += [activition_function()]
        self.mlp = nn.Sequential(*layers)

    def forward(self, sampling_locations):
        return self.mlp(sampling_locations)

class nerf_mlp(nn.Module):
    def __init__(self):
        super().__init__()

        self.sigma_mlp = ClassicalMLP(3,
                                      256 + 1,
                                      256,
                                      5)

        self.rgb_mlp = ClassicalMLP(256 + 3,
                                    3,
                                    256,
                                    5)


    def forward(self, sampling_locations, directions):
        sigma_plus_feature_vector = self.sigma_mlp.forward(sampling_locations)
        sigma_values = (sigma_plus_feature_vector[..., 0]).abs()
        feature_vector = sigma_plus_feature_vector[..., 1:]
        rgb_input = torch.cat([feature_vector, directions], dim=-1)
        rgb_values = torch.tanh(self.rgb_mlp.forward(rgb_input).abs())
        return_values = torch.cat([sigma_values.unsqueeze(-1), rgb_values], dim=-1)
        return return_values

class Nerf3DModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.integrator = NerfIntegrator()
        self.coarse_mlp = nerf_mlp()
        self.fine_mlp = nerf_mlp()

    def evaluate_color_from_rays(self, network_values_at_samples, interval_sizes):
        rgb_values = network_values_at_samples[..., 1:4]
        sigma_values = network_values_at_samples[..., 0]
        colors = self.integrator(rgb_values, sigma_values, interval_sizes)
        return colors, sigma_values

    def forward(self, input_x):
        rays = input_x["ray_directions"]
        camera_positions = input_x["camera_origin"]
        rays = rays / rays.norm(dim=-1, keepdim=True)
# i replaced sampler netwroks with tensors of ones, to simplify the code
# the problem presists

        shape = list(rays.shape)
        shape = shape[:-1] + [64] + [shape[-1]]
        coarse_sampling_locations = torch.ones(shape, device=rays.device)
        sampling_start = torch.ones(shape, device=rays.device)
        sampling_end = torch.ones(shape, device=rays.device)
        coarse_interval_sizes = torch.ones(shape[:-1], device=rays.device)

        repeats = [1] * (len(coarse_sampling_locations.shape) - 2) + [coarse_sampling_locations.shape[-2]] + [1]
        coarse_directions = rays.unsqueeze(-2).repeat(repeats)
        network_values_at_samples = self.coarse_mlp(coarse_sampling_locations, coarse_directions)
        # the first 3 outputs of the DNN are the rgb channels and the last one is
        # the sigma channel
        resulting_coarse_image, sigma_values = self.evaluate_color_from_rays(network_values_at_samples,
                                                                             coarse_interval_sizes)

        shape[-2] = 192
        fine_sampling_points = torch.ones(shape, device=rays.device)
        fine_interval_sizes = torch.ones(shape[:-1], device=rays.device)

        repeats = [1]*(len(rays.shape) - 1) + [fine_sampling_points.shape[-2]] + [1]
        fine_directions = rays.unsqueeze(-2).repeat(repeats)
        network_values_at_fine_samples = self.fine_mlp(fine_sampling_points, fine_directions)
        resulting_fine_image, _ = self.evaluate_color_from_rays(network_values_at_fine_samples, fine_interval_sizes)

        network_output = {
            "rgb_coarse": resulting_coarse_image,
            "rgb_fine": resulting_fine_image,
            "rays": rays,
            "camera_positions": camera_positions,
            "sampling_locations_coarse": coarse_sampling_locations,
            "fine_sampling_locations": fine_sampling_points,
        }
        return network_output

class LossFunction(nn.Module):
    def forward(self, input_batch, prediction):
        gb_rgb = input_batch["rgb"]
        predication_rgb_coarse = prediction["rgb_coarse"]
        predication_rgb_fine = prediction["rgb_fine"]
        coarse_rgb_loss = (gb_rgb - predication_rgb_coarse).norm(dim=-1).mean()
        fine_rgb_loss = (gb_rgb - predication_rgb_fine).norm(dim=-1).mean()
        return_value = coarse_rgb_loss + fine_rgb_loss
        return {"loss_value": return_value, "coarse_rgb_loss":coarse_rgb_loss, "fine_rgb_loss":fine_rgb_loss}

class RandomDataset(Dataset):
    def __init__(self):
        super().__init__()
        self.data = [{
            "rgb": torch.rand(800, 800, 3).cuda().reshape(-1, 3),
            "ray_directions": torch.rand(800, 800, 3).cuda().reshape(-1, 3),
            "camera_origin": torch.rand(1, 3).cuda(),
        } for _ in range(10)]
        self.sample_length = self.data[0]["rgb"].shape[0]
        self.chunk_size = 1024
        self.chunks_per_sample = int(np.ceil(self.sample_length / self.chunk_size))
        self.length = self.chunks_per_sample * len(self.data)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        sample_index = idx // self.chunks_per_sample
        batch_index = idx % self.chunks_per_sample
        sample = self.data[sample_index]
        start = batch_index * self.chunk_size
        end = (batch_index + 1) * self.chunk_size
        end = min(end, self.sample_length)
        pixel_keys = ["rgb","ray_directions"]
        chunk = {
            key: sample[key][start:end] for key in pixel_keys
        }
        chunk["camera_origin"] = sample["camera_origin"]
        return chunk


if __name__ == '__main__':

    dataset = RandomDataset()
    data_loader = DataLoader(dataset, batch_size=1)

    device = torch.device("cuda")
    model = Nerf3DModel().to(device)
    loss_function = LossFunction()
    optimizer = torch.optim.Adam(model.parameters())
    for epoch in range(1000):
        for batch in data_loader:
            t1 = time()
            optimizer.zero_grad()
            pred = model(batch)
            loss = loss_function(batch,pred)
            loss["loss_value"].backward()
            optimizer.step()
            print(time() - t1)


this coda example shows the problem.
i am running on GTX 3090. so if there is not enough sapce on the gpu, you can change
“chunk size” in the dataset to smothing smaller then 1024
the code printed:
0.6795637607574463
0.01598954200744629
0.007995367050170898
0.01598811149597168
0.00799417495727539
0.10390162467956543
0.19186925888061523
0.17587995529174805
0.18573331832885742
0.16800951957702637
0.19200444221496582
0.16811847686767578
0.19177031517028809
0.1763777732849121
as you can see, the first iterations are fast, but then come the slowdown