Implementing custom loss function

Hi, I am trying to implement a regression model that will predict the absolute angle of rotation of an image. For this, I need to implement a custom loss function, however I cannot figure out how to do this.

I first prototyped what this function would look like for a single training instance. It finds the angle between y and yhat:
def loss(y, yhat):
option1 = abs(y-yhat)
option2 = 360 - option1
return min(option1,option2)

However, I am lost when it comes to how I would define this function to be compatible with fastai. Based on my research, it seems like I would need to wrap my function in the BaseLoss class. Could someone show me what this is supposed to look like? I couldn’t find examples online.

I’ve seen that once I figure this out, I can just pass it to the learner:
learn.loss_func = my_loss_function
Then I can train just like normal.

Note: I am working with fastai v2

Just pass loss_func=loss I believe. If tensor types throw an issue for you (thanks PyTorch) then you would just pass your func to BaseLoss like so:

loss_func = BaseLoss(loss)

So a full setup for you to try:

learn = cnn_learner(net, loss_func=loss) #this can be any _learner or Learner

Or:

loss_func = BaseLoss(loss)
learn = cnn_learner(net, loss_func=loss_func)

The other option if you just want to pass in loss_func=loss is cast your tensors to TensorBase when they first get sent to your loss function, so something like:

def loss(y, yhat):
  y,yhat = TensorBase(y), TensorBase(yhat)
  option1 = abs(y-yhat)
  option2 = 360 - option1
  return min(option1,option2)