Class: Gsplat::Training::Losses::RegularizedReconstruction

Inherits:
Autograd::Function show all
Defined in:
lib/gsplat/training/losses.rb

Overview

Reconstruction loss with activated opacity and scale mean penalties.

Class Method Summary collapse

Methods inherited from Autograd::Function

apply

Class Method Details

.backward(context, grad_output) ⇒ Object

rubocop:disable Metrics/AbcSize



103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
 
# File 'lib/gsplat/training/losses.rb', line 103

def backward(context, grad_output)
  difference, count, cache, ssim_lambda,
    opacity_shape, scale_shape, opacity_reg, scale_reg = context.saved_values
  signs = difference.class.zeros(*difference.shape)
  signs[difference.gt(0)] = 1
  signs[difference.lt(0)] = -1
  l1_gradient = signs * ((1 - ssim_lambda) * grad_output.to_f / count)
  ssim_a, ssim_b = Math::Ssim.backward(cache, grad_output)
  opacity_gradient = difference.class.ones(*opacity_shape) *
                     opacity_reg * grad_output.to_f / opacity_shape.inject(:*)
  scale_gradient = difference.class.ones(*scale_shape) *
                   scale_reg * grad_output.to_f / scale_shape.inject(:*)
  [
    l1_gradient - (ssim_lambda * ssim_a),
    -l1_gradient - (ssim_lambda * ssim_b),
    opacity_gradient,
    scale_gradient
  ]
end

.forward(context, prediction, target, opacities, scales, ssim_lambda:, opacity_reg:, scale_reg:, layout:) ⇒ Object

rubocop:disable Metrics/ParameterLists



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

def forward(context, prediction, target, opacities, scales, ssim_lambda:, opacity_reg:, scale_reg:, layout:)
  # rubocop:enable Metrics/ParameterLists
  difference = prediction - target
  score, ssim_cache = Math::Ssim.forward(prediction, target, layout: layout)
  context.save(
    difference, prediction.size, ssim_cache, ssim_lambda,
    opacities.shape, scales.shape, opacity_reg, scale_reg
  )
  reconstruction = ((1 - ssim_lambda) * difference.abs.mean.to_f) + (ssim_lambda * (1 - score))
  prediction.class.cast(
    reconstruction + (opacity_reg * opacities.mean.to_f) + (scale_reg * scales.mean.to_f)
  )
end