Custom likelihood and data_loader

[izzatum]

Hi! Is there any way that we can implement custom likelihood instead of 'regression' and 'classification', and data_loader? I'm trying to use laplace for PINN. So, the negative log-likelihood (loss) and data_loader are slightly different.

My PINN network has two inputs. I faced this issue:

`/usr/local/lib/python3.7/dist-packages/laplace/baselaplace.py in fit(self, trainloader) 120 self.model.eval() 121 --> 122 X, = next(iter(train_loader)) 123 with torch.no_grad(): 124 self.n_outputs = self.model(X[:1].to(self._device)).shape[-1]

ValueError: too many values to unpack (expected 2)`

What would you advise in this case? Thanks!

[edaxberger]

Hi, thanks for reaching out, and sorry for the late response (I somehow thought someone had answered already)!

Could you provide more details on the exact likelihood that you want to use?

Concerning the issue with the data loader, you could just re-define your data set/loader to only produce 2 objects (i.e. the inputs X and the labels y), which is the convention we use in the code.

Hope this helps, otherwise please let us know!

[runame]

Concerning the data loader, it also needs to fulfill some other requirements:

• data_loader.dataset needs to exist and have __len__() implemented.
• The inputs X need to have a shape and device attribute, have __len__() implemented, and need to accept slices, e.g. X[:1].

I will try to eliminate some of the requirements mentioned above. It would be very helpful if you could give us feedback on which other parts of your data loader cause problems. Thanks!

[izzatum]

Hi Erik and Runa,

Great thanks for your kind assistance!

Regarding the loss/likelihood, the physics-informed neural network has multiple loss functions, yet all of them are just MSELoss. The issue here is that to compute that loss I need to take the derivatives of the network's output. So, it is not straightforward. Kindly find my implementation below for the loss function.

For the data loader, it depends on the implementation. Y is not needed, and X is multiple columns that represent the multiple input of the network. As in my implementation below, I just made Y as a dummy variable.

I hope this would help.

Thanks a lot!

`class LossEikonal(nn.Module): """ Create a loss function for Eikonal equation physics informed neural network (PINN).

Attributes: model: PINN model T0: smooth background travel time. px0: smooth background travel time derivative w.r.t x pz0: smooth background travel time derivative w.r.t z vel: smooth background velocity model sids: source indices """'

'def init(self, model, params): super().init() self.model = model self.T0 = params['T0'] self.px0 = params['px0'] self.pz0 = params['pz0'] self.vel = params['vel'] self.sids = params['sids']'

'def forward(self, dataset: Tensor, dummy=None) -> Tensor: """ Compute loss function for Eikonal equation physics informed neural network (PINN).

Args:
  tau: the PINN predictions
  dataset: current dataset

Returns:
  loss: Tensor of loss value
"""
# number of data points
N = dataset.shape[0]

# compute derivative of tau
tau, dtau_dz, dtau_dx = tau_derivative(self.model, dataset)

# compute factored eikonal loss
loss_eqn = l2(
              (self.T0*dtau_dx + tau*self.px0)**2 + 
              (self.T0*dtau_dz + tau*self.pz0)**2 - 
              1.0/self.vel**2
            )**2/N

# compute eikonal equation solution constraint with Heaviside function
loss_heav = l2((1 - torch.sign(tau*self.T0))*torch.abs(
                tau*self.T0))**2/N

# compute source constraint for eikonal equation
loss_source = l2(tau[self.sids] - 1.)**2

return loss_eqn + loss_heav + loss_source'

'@staticmethod def tau_derivative(model, dataset): """ Compute derivative of tau (PINN output) w.r.t data

Args:
  model: PINN model
  dataset: data

Returns:
  dtau_dz: derivative of tau w.r.t z
  dtau_dx: derivative of tau w.r.t x
"""

# enable grad for dataset
#dataset.requires_grad = True

# compute gradient of tau w.r.t dataset
tau = model(dataset)

with torch.no_grad():
  grad_multiplier = torch.ones_like(tau, requires_grad=False).to(device)
  g = ag.grad(outputs=tau, inputs=dataset, grad_outputs=grad_multiplier,
              create_graph=True)[0]

  dtau_dz = g[..., 0].reshape(-1, 1)
  dtau_dx = g[..., 1].reshape(-1, 1)

tau.detach()

return tau, dtau_dz, dtau_dx`

[wiseodd]

Unfortunately, we don't currently support custom loss functions. We currently only supports CrossEntropyLoss and MSELoss since those are the losses supported by the Hessian backends we use (BackPACK and ASDL).