“Compute gradient penalty: (L2_norm(dy/dx))**2.”

  1. class R1_reg(nn.Module):
  2.     def __init__(self, lambda_r1=10.0):
  3.         super(R1_reg, self).__init__()
  4.         self.lambda_r1 = lambda_r1
  6.     def __call__(self, d_out, d_in):
  7.         """Compute gradient penalty: (L2_norm(dy/dx))**2."""
  8.         b = d_in.shape[0]
  9.         dydx = torch.autograd.grad(outputs=d_out.mean(),
  10.                                    inputs=d_in,
  11.                                    retain_graph=True,
  12.                                    create_graph=True,
  13.                                    only_inputs=True)[0]
  14.         dydx_sq = dydx.pow(2)
  15.         assert (dydx_sq.size() == d_in.size())
  16.         r1_reg = dydx_sq.sum() / b
  18.         return r1_reg * self.lambda_r1

SelectiveClassesNonSatGANLoss

  2. class SelectiveClassesNonSatGANLoss(nn.Module):
  3.     def __init__(self):
  4.         super(SelectiveClassesNonSatGANLoss, self).__init__()
  5.         self.sofplus = nn.Softplus()
  7.     def __call__(self, input, target_classes, target_is_real, is_gen=False):
  8.         bSize = input.shape[0]
  9.         b_ind = torch.arange(bSize).long()
  10.         relevant_inputs = input[b_ind, target_classes, :, :]
  11.         if target_is_real:
  12.             loss = self.sofplus(-relevant_inputs).mean()
  13.         else:
  14.             loss = self.sofplus(relevant_inputs).mean()
  16.         return loss