Skip to content

train

train #

train — a splat from posed images, on the GPU one call at a time.

  1. Toolchain. The trainer runs in the interpreter CG_GPU_PYTHON names (else the one ns-train on the path belongs to): torch, gsplat and nerfstudio live there, never in the library's own environment.
  2. Lock. The whole run holds cg.lock("gpu"), so calls from every page worker, process and script on the machine train one after another and a waiting call says so in its progress.
  3. Progress. The driver writes progress.json as it trains; a watcher thread reads it into a Training step, so the count, loss and gaussian count reach a page or terminal while the run is live.

GPU_PYTHON_ENV = 'CG_GPU_PYTHON' module-attribute #

Names the interpreter holding torch, gsplat and nerfstudio.

gpu_python() #

The interpreter the trainer runs in.

Raises:

Type Description
RecipeError

If CG_GPU_PYTHON is unset and ns-train is not on the path.
Source code in capturegraph-lib/capturegraph/recipes/splats/train.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def gpu_python() -> Path:
    """The interpreter the trainer runs in.

    Raises:
        RecipeError: If ``CG_GPU_PYTHON`` is unset and ``ns-train`` is not on the path.
    """
    named = os.environ.get(GPU_PYTHON_ENV)
    if named:
        return Path(named)
    ns_train = shutil.which("ns-train")
    if ns_train is not None:
        return Path(ns_train).with_name("python")
    raise RecipeError(
        f"No GPU interpreter for splat training: set {GPU_PYTHON_ENV} to a python with "
        f"torch, gsplat and nerfstudio installed, or put its ns-train on the PATH."
    )

train(frames, seed, iterations=15000, optimize_cameras=False) #

A 3D Gaussian splat trained on frames from seed with nerfstudio's splatfacto.

frames = cg.splats.scene_dataset(session.capture)
splat = cg.splats.train(frames, cg.splats.seed(frames))

Parameters:

Name Type Description Default
frames Array[PosedImage]

The posed images to fit.

 required
seed PointCloud

The coloured points the gaussians start from.

 required
iterations int

How many optimizer steps to run.

 15000
optimize_cameras bool

Whether to refine the camera poses while training, for poses a sensor estimated rather than a tracker.

 False

Raises:

Type Description
RecipeError

If no GPU interpreter can be found.
CalledProcessError

If training fails; the exception carries the tail of the trainer's log.
Source code in capturegraph-lib/capturegraph/recipes/splats/train.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@pure_function("splats.train/1")
def train(
    frames: Array[PosedImage],
    seed: PointCloud,
    iterations: int = 15_000,
    optimize_cameras: bool = False,
) -> Splat:
    """A 3D Gaussian splat trained on ``frames`` from ``seed`` with nerfstudio's splatfacto.

    ```python
    frames = cg.splats.scene_dataset(session.capture)
    splat = cg.splats.train(frames, cg.splats.seed(frames))
    ```

    Args:
        frames: The posed images to fit.
        seed: The coloured points the gaussians start from.
        iterations: How many optimizer steps to run.
        optimize_cameras: Whether to refine the camera poses while training,
            for poses a sensor estimated rather than a tracker.

    Raises:
        RecipeError: If no GPU interpreter can be found.
        subprocess.CalledProcessError: If training fails; the exception carries
            the tail of the trainer's log.
    """
    work = ManagedDir.new()
    splat = Splat.new("ply")
    command = [
        gpu_python(),
        TRAINER,
        "--frames",
        _stored(frames, work),
        "--output",
        work,
        "--iterations",
        str(iterations),
        "--seed",
        seed,
        "--splat",
        splat,
        "--progress",
        work / PROGRESS,
        *(["--optimize-cameras"] if optimize_cameras else []),
    ]
    with lock(GPU), step("Training", total=iterations) as training:
        with _watching(training, work / PROGRESS):
            work.exec(command, env={**os.environ, **_TORCH_ENV})
    return splat