Class: Gsplat::Training::Trainer

Inherits:
Object
  • Object
show all
Defined in:
lib/gsplat/training/trainer.rb

Overview

End-to-end multi-view Gaussian training loop.

Defined Under Namespace

Classes: Result

Constant Summary collapse

LEARNING_RATES =

Mapping from parameter names to Config learning-rate attributes.

{
  means: :means_lr,
  scales: :scales_lr,
  quats: :quats_lr,
  opacities: :opacities_lr,
  sh0: :sh0_lr,
  shN: :shN_lr
}.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(scene:, config: Config.new, params: nil, strategy: nil) ⇒ Trainer

Returns a new instance of Trainer.



23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
 
# File 'lib/gsplat/training/trainer.rb', line 23

def initialize(scene:, config: Config.new, params: nil, strategy: nil)
  @scene = scene
  @config = config
  @rng = Random.new(config.seed)
  @params = params || Utils.init_from_points(
    scene.points,
    scene.colors,
    sh_degree: config.sh_degree,
    init_opacity: config.init_opacity,
    init_scale: config.init_scale,
    rng: @rng
  )
  @optimizers = build_optimizers
  @strategy = strategy || Strategy::Default.new
  @strategy.check_sanity(@params, @optimizers)
  @strategy_state = @strategy.initialize_state(scene_scale: scene.scene_scale)
  @scheduler = Optim::ExponentialLR.new(
    @optimizers.fetch(:means),
    lr_final: config.means_lr_final * scene.scene_scale,
    max_steps: config.max_steps
  )
  @step = 0
end

Instance Attribute Details

#config ⇒ Object (readonly)

Returns the value of attribute config.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def config
  @config
end

#optimizers ⇒ Object (readonly)

Returns the value of attribute optimizers.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def optimizers
  @optimizers
end

#params ⇒ Object (readonly)

Returns the value of attribute params.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def params
  @params
end

#scene ⇒ Object (readonly)

Returns the value of attribute scene.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def scene
  @scene
end

#step ⇒ Object (readonly)

Returns the value of attribute step.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def step
  @step
end

#strategy ⇒ Object (readonly)

Returns the value of attribute strategy.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def strategy
  @strategy
end

#strategy_state ⇒ Object (readonly)

Returns the value of attribute strategy_state.



21
22
23
 
# File 'lib/gsplat/training/trainer.rb', line 21

def strategy_state
  @strategy_state
end

Instance Method Details

#evaluate ⇒ Hash{Symbol=>Float}

Renders every view without graph recording and reports PSNR/SSIM.

Returns:

  • (Hash{Symbol=>Float}) 


73
74
75
76
77
78
79
80
81
82
 
# File 'lib/gsplat/training/trainer.rb', line 73

def evaluate
  Autograd.no_grad do
    rendered, = render((0...scene.camera_count).to_a, sh_degree_for(step))
    values = Ops::TensorOps.data(rendered)
    {
      psnr: Losses.psnr(values, scene.images),
      ssim: Losses.ssim(values, scene.images, layout: :nhwc)
    }
  end
end

#load_checkpoint!(source) ⇒ IO::Checkpoint::Snapshot

Restores parameters, optimizer moments, step, and scheduler position.

Parameters:

  • source (String, #read)

Returns:



105
106
107
108
109
110
 
# File 'lib/gsplat/training/trainer.rb', line 105

def load_checkpoint!(source)
  snapshot = IO::Checkpoint.restore!(source, params: params, optimizers: optimizers)
  @step = snapshot.step
  @scheduler.step(step)
  snapshot
end

#save_checkpoint!(path = nil) ⇒ String

Saves current parameters, optimizer moments, step, and config.

Parameters:

  • path (String, nil) (defaults to: nil) —

    generated from config.output_dir when nil

Returns:

  • (String) —

    written path



88
89
90
91
92
93
94
95
96
97
98
99
 
# File 'lib/gsplat/training/trainer.rb', line 88

def save_checkpoint!(path = nil)
  path ||= File.join(config.output_dir, "checkpoints", format("step_%06d.npz", step))
  FileUtils.mkdir_p(File.dirname(path))
  IO::Checkpoint.save(
    path,
    params: params,
    optimizers: optimizers,
    step: step,
    config: config.to_h
  )
  path
end

#train(steps: nil) ⇒ Result

Optimizes until the configured maximum or for an additional step count.

Parameters:

  • steps (Integer, nil) (defaults to: nil)

Returns:



51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
 
# File 'lib/gsplat/training/trainer.rb', line 51

def train(steps: nil)
  target_step = steps ? step + steps : config.max_steps
  initial_metrics = evaluate
  history = []
  while step < target_step
    @step += 1
    loss = training_step
    history << { step: step, loss: loss } if log_due?
    evaluate_and_log!(history) if config.eval_steps.include?(step)
    save_checkpoint! if config.save_steps.include?(step)
  end
  Result.new(
    step: step,
    initial_metrics: initial_metrics,
    final_metrics: evaluate,
    history: history.freeze
  )
end