Class: Gsplat::Training::Losses::Reconstruction

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

Overview

Weighted L1 plus (1 - SSIM) reconstruction objective.

Class Method Summary collapse

Methods inherited from Autograd::Function

apply

Class Method Details

.backward(context, grad_output) ⇒ Object

This method is part of a private API. You should avoid using this method if possible, as it may be removed or be changed in the future.

Propagates the combined reconstruction gradient.



72
73
74
75
76
77
78
79
80
 
# File 'lib/gsplat/training/losses.rb', line 72

def backward(context, grad_output)
  difference, count, ssim_cache, ssim_lambda = 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(ssim_cache, grad_output)
  [l1_gradient - (ssim_lambda * ssim_a), -l1_gradient - (ssim_lambda * ssim_b)]
end

.forward(context, prediction, target, ssim_lambda:, layout:) ⇒ Object

This method is part of a private API. You should avoid using this method if possible, as it may be removed or be changed in the future.

Evaluates weighted L1 and SSIM reconstruction loss.



59
60
61
62
63
64
65
66
67
68
 
# File 'lib/gsplat/training/losses.rb', line 59

def forward(context, prediction, target, ssim_lambda:, layout:)
  unless prediction.shape == target.shape
    raise ShapeError, "reconstruction shapes differ: #{prediction.shape.inspect} and #{target.shape.inspect}"
  end

  difference = prediction - target
  score, ssim_cache = Math::Ssim.forward(prediction, target, layout: layout)
  context.save(difference, prediction.size, ssim_cache, ssim_lambda)
  prediction.class.cast(((1 - ssim_lambda) * difference.abs.mean.to_f) + (ssim_lambda * (1 - score)))
end