Train with "image tiles" to save VRAM.

[ArthurBrussee]

To reduce VRAM it should be possible to train on images of higher resolution in chunks, by eg. splitting the image in 4 quadrants and using each as a seperate training sample.

Likely would go great with #38 , if implemented.

This would still keep quite a few buffers at full size, so it's not like training on 1/4th of an image would save 75% of VRAM, but it would be significant.

[yuyou]

Claude suggested a plan, check if it makes sense or correct:

Plan: N×N Image Grid Splitting for Training

Implement configurable image splitting into an N×N grid (e.g., 2×2 = 4 quadrants, 3×3 = 9 tiles) for training, processing each tile sequentially as a separate batch with adjusted camera parameters. This enables training on higher resolution images with lower memory usage and better detail preservation.

Steps

1. Add configuration option

File: crates/brush-dataset/src/config.rs (lines 15-36)

• Add split_image_grid: Option<u32> field to LoadDataseConfig struct
• Add CLI argument with #[arg(long, help_heading = "Dataset Options")]
• Document that N=2 creates 2×2 grid (4 tiles), N=3 creates 3×3 grid (9 tiles), etc.
• Default to None (disabled)

2. Create helper function for camera adjustment

File: crates/brush-dataset/src/scene.rs

• Add adjust_camera_for_tile(camera: &Camera, grid_size: u32, tile_x: u32, tile_y: u32) -> Camera
• Formula for each tile at position (tile_x, tile_y) in an N×N grid:
  • new_center_uv.x = N * camera.center_uv.x - tile_x
  • new_center_uv.y = N * camera.center_uv.y - tile_y
• FOV and focal length remain unchanged (we're cropping, not zooming)
• Position and rotation unchanged (same viewpoint)

3. Create image cropping function

File: crates/brush-dataset/src/scene.rs

• Add crop_image_tile(image: &DynamicImage, grid_size: u32, tile_x: u32, tile_y: u32) -> DynamicImage
• Calculate tile dimensions: tile_w = width / grid_size, tile_h = height / grid_size
• Crop coordinates: x = tile_x * tile_w, y = tile_y * tile_h
• Use image.crop_imm(x, y, tile_w, tile_h) to extract the tile
• Handle alpha channels correctly (preserve from original)

4. Modify batch generation

File: crates/brush-dataset/src/scene_loader.rs (lines 81-125)

• In the main loop (line ~82), after selecting a view, check if grid splitting is enabled
• If grid_size is Some(N), generate N² batches instead of 1:
  • Nested loop: for tile_y in 0..N { for tile_x in 0..N { ... } }
  • For each tile: crop image → adjust camera → create batch
  • Send batches sequentially to maintain order
• If grid_size is None, use existing single-batch logic
• Process tiles in row-major order (left→right, top→bottom)

5. Update Scene/SceneView structures

File: crates/brush-dataset/src/scene.rs

• Pass split_image_grid config through to SceneLoader::new()
• Modify SceneLoader struct to store grid_size: Option<u32>
• No changes needed to SceneView or Scene themselves (splitting happens at batch time)

6. Handle image cache correctly

File: crates/brush-dataset/src/scene_loader.rs (lines 32-50)

• Cache the full original image, not individual tiles (more memory efficient)
• After retrieving from cache, crop to tiles on-demand
• ImageCache semantics remain unchanged (caches by view index)

7. Update training loop context

File: crates/brush-process/src/train_stream.rs (line 185)

• No code changes needed - next_batch() already yields batches sequentially
• Training metrics will naturally reflect tile-based training
• Each iteration processes one tile (or full image if splitting disabled)

8. Pass config through dataset loading

Files:

• crates/brush-dataset/src/formats/mod.rs

• crates/brush-dataset/src/formats/colmap.rs

• crates/brush-dataset/src/formats/nerfstudio.rs

• Verify load_args: &LoadDataseConfig is passed to Scene creation

• Pass split_image_grid config to SceneLoader::new() call

• Both train and eval datasets should use the same grid setting

Verification

• Compile: cargo build --release
• Test with 2×2 grid: ./target/release/brush /path/to/dataset --split-image-grid 2 --eval-split-every 8
• Check logs show correct batch count: Should see 4× more iterations per epoch
• Test with 3×3 grid: --split-image-grid 3 (9 tiles per image)
• Verify camera parameters: Add debug logging in adjusted camera to confirm principal point shifts
• Test memory usage: Should be lower for same source resolution
• Compare final quality: Train with/without splitting to verify quality parity

Decisions

• Sequential batching: Process all N² tiles from one image before moving to next, maintaining consistent view coverage
• Consistent eval: Both training and evaluation use the same grid splitting to ensure fair metrics
• Flexible N×N: Support arbitrary grid sizes (2×2, 3×3, 4×4, etc.) via --split-image-grid N parameter, not hardcoded to quadrants
• Batch-time splitting: Split images during batch generation in SceneLoader rather than preprocessing, preserving original dataset structure and cache efficiency

Camera Mathematics Explanation

When splitting an image into N×N tiles, the camera's intrinsic parameters must be adjusted for each tile:

Principal Point (center_uv):

• Original image: principal point is normalized (e.g., 0.5, 0.5 for center)
• For tile at position (tile_x, tile_y):
  • The original principal point in pixel coordinates: px = center_uv.x * width
  • After cropping to tile: new_px = px - tile_x * (width / N)
  • Normalized to tile size: new_center_uv.x = new_px / (width / N) = N * center_uv.x - tile_x
  • Same logic for y-axis

Field of View (FOV):

• FOV is determined by the focal length and sensor size
• When cropping, we're not changing the lens, so FOV angles remain the same
• The rendering will naturally show the cropped portion based on the adjusted principal point

Position & Rotation:

• Camera position and orientation in 3D space remain unchanged
• Each tile is rendered from the exact same viewpoint, just looking at a different portion of the image plane