code: # Copyright (c) 2018, Curious AI Ltd. All rights reserved.
  #
  # This work is licensed under the Creative Commons Attribution-NonCommercial
  # 4.0 International License. To view a copy of this license, visit
  # http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
  # Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
  
  """Custom loss functions"""
  
  import torch
  from torch.nn import functional as F
  from torch.autograd import Variable
  import pdb
  import numpy as np
  
  def softmax_mse_loss(input_logits, target_logits):
      """Takes softmax on both sides and returns MSE loss
  
      Note:
      - Returns the sum over all examples. Divide by num_classes
        Divide by the batch size afterwards if you want the mean.
      - Sends gradients to inputs but not the targets.
      """
      assert input_logits.size() == target_logits.size()
      input_softmax = F.softmax(input_logits, dim=1)
      target_softmax = F.softmax(target_logits, dim=1)
      num_classes = input_logits.size()[1]
      return F.mse_loss(input_softmax, target_softmax, size_average=False) / num_classes
  
  
  def softmax_kl_loss(input_logits, target_logits):
      """Takes softmax on both sides and returns KL divergence
  
      Note:
      - Returns the sum over all examples. Divide by the batch size afterwards
        if you want the mean.
      - Sends gradients to inputs but not the targets.
      """
      assert input_logits.size() == target_logits.size()
      input_log_softmax = F.log_softmax(input_logits, dim=1)  # log(q)
      target_softmax = F.softmax(target_logits, dim=1) # p
      return F.kl_div(input_log_softmax, target_softmax, size_average=False)
  
  
  def symmetric_mse_loss(input1, input2):
      """Like F.mse_loss but sends gradients to both directions.
          cuz input1/input2 are tensors with grad, while target in F.mse_loss has no grad.
  
      Note:
      - Returns the sum over all examples. Divide by the batch size afterwards
        if you want the mean.
      - Sends gradients to both input1 and input2.
      """
      assert input1.size() == input2.size()
      num_classes = input1.size()[1]
      return torch.sum((input1 - input2)**2) / num_classes
apis: torch.sum torch.nn.functional.log_softmax torch.nn.functional.mse_loss torch.nn.functional.softmax torch.nn.functional.kl_div