text: 
  
  
  
  
  from caffe2.python import core, workspace
  from hypothesis import given, settings
  import caffe2.python.hypothesis_test_util as hu
  import caffe2.python.serialized_test.serialized_test_util as serial
  import hypothesis.strategies as st
  import numpy as np
  
  import unittest
  
  
  class TestSoftmaxOps(serial.SerializedTestCase):
  
      @serial.given(n=st.sampled_from([0, 2, 4, 71, 103]),
                    D=st.sampled_from([0, 4, 8, 64, 79, 256, 333]),
                    engine=st.sampled_from([None, 'CUDNN']),
                    **hu.gcs)
      def test_softmax(self, n, D, engine, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          X = np.random.rand(n, D).astype(np.float32)
          X = X + 1e-2
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax(X):
              probs = np.zeros((n, D))
              rowmax = np.zeros(n)
  
              if D == 0:
                  return [probs]
  
              for i in range(n):
                  rowmax[i] = max(X[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
  
              return [probs]
  
          op = core.CreateOperator(
              "Softmax",
              ["X"],
              ["probs"],
              engine=engine
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X],
              reference=label_softmax,
          )
  
      @given(n=st.sampled_from([0, 2, 4, 71, 103, 555, 751, 1201]),
                    D=st.sampled_from([0, 4, 8, 64, 79, 256, 333, 1000]),
                    engine=st.sampled_from([None, 'CUDNN']),
                    **hu.gcs)
      @settings(deadline=10000)
      def test_softmax_grad(self, n, D, engine, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          Y = np.random.rand(n, D).astype(np.float32)
          dY = np.random.rand(n, D).astype(np.float32)
          Y = Y + 1e-2
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax_grad(X, dY):
              dX = Y * 0.0
              for i in range(n):
                  d = np.dot(Y[i, :], dY[i, :])
                  dX[i, :] = Y[i, :] * (dY[i, :] - d)
              return [dX]
  
          op = core.CreateOperator(
              "SoftmaxGradient",
              ["Y", "dY"],
              ["dX"],
              engine=engine
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[Y, dY],
              reference=label_softmax_grad,
          )
  
      @given(axis=st.integers(min_value=1, max_value=4),
             engine=st.sampled_from([None, 'CUDNN']),
             **hu.gcs)
      def test_softmax_axis(self, axis, engine, gc, dc):
          np.random.seed(1)
          X = np.random.randn(1, 2, 3, 2, 1).astype(np.float32)
          X = X + 1e-2
  
          def prod(xs):
              p = 1
              for x in xs:
                  p *= x
              return p
  
          N = prod(list(X.shape)[:axis])
          D = prod(list(X.shape)[axis:])
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax(X):
              X_ = X.reshape(N, D)
              probs = np.zeros((N, D))
              rowmax = np.zeros(N)
              for i in range(N):
                  rowmax[i] = max(X_[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X_[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
  
              return [probs.reshape(*X.shape)]
  
          op = core.CreateOperator(
              "Softmax",
              ["X"],
              ["probs"],
              axis=axis,
              engine=engine,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X],
              reference=label_softmax,
          )
  
          self.assertGradientChecks(
              gc, op, [X], 0, [0], stepsize=1e-4, threshold=1e-2)
  
      @given(n=st.integers(2, 10), D=st.integers(4, 16),
             only_loss=st.booleans(), **hu.gcs)
      @settings(deadline=10000)
      def test_softmax_with_loss(self, n, D, gc, only_loss, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          np.random.seed(2603)
          X = np.random.rand(n, D).astype(np.float32)
          X = X + 1e-2
  
          # Initialize label
          label = (np.random.rand(n) * D).astype(np.int32)
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax_crossent(X, label):
              probs = np.zeros((n, D))
              rowmax = np.zeros(n)
              for i in range(n):
                  rowmax[i] = max(X[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
  
              label_xent = [-np.log(max(probs[i][label[i]], 1e-20))
                            for i in range(n)]
              avgloss = np.sum(label_xent) / float(n)
              return (probs, avgloss)
  
          op = core.CreateOperator(
              "SoftmaxWithLoss",
              ["X", "label"],
              ["probs", "avgloss"],
              only_loss=only_loss,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X, label],
              reference=label_softmax_crossent,
          )
  
          self.assertGradientChecks(
              gc, op, [X, label], 0, [1], stepsize=1e-4, threshold=1e-2)
  
      @given(
          n=st.integers(2, 5),
          D=st.integers(4, 16),
          only_loss=st.booleans(),
          label_prob=st.booleans(),
          **hu.gcs
      )
      @settings(deadline=10000)
      def test_softmax_with_loss_axis_2(
          self, n, D, only_loss, label_prob,
          gc, dc
      ):
          np.random.seed(2603)
          X = np.random.rand(n, n, D).astype(np.float32)
          X = X + 1e-2
  
          if label_prob:
              label = np.random.rand(n, n, D).astype(np.float32)
              label /= label.sum(axis=2, keepdims=True)
          else:
              label = (np.random.rand(n, n) * D).astype(np.int32)
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax_crossent(X, label):
              probs = np.zeros((n, n, D))
              rowmax = np.zeros((n, n))
              for i in range(n):
                  for j in range(n):
                      rowmax[i, j] = max(X[i, j, ])
                      # We need to subtract the max to avoid numerical issues
                      probs[i, j] = X[i, j] - rowmax[i, j]
                      exps = np.exp(probs[i, j, ])
                      norm = sum(exps)
                      probs[i, j, ] = exps / norm
              label_xent = 0
              for i in range(n):
                  for j in range(n):
                      if label_prob:
                          for k in range(D):
                              label_xent += (
                                  -np.log(max(probs[i, j, k], 1e-20)) *
                                  label[i, j, k]
                              )
                      else:
                          label_xent += -np.log(max(probs[i, j, label[i, j]], 1e-20))
  
              avgloss = label_xent / float(n * n)
              return (probs, avgloss)
  
          op = core.CreateOperator(
              "SoftmaxWithLoss",
              ["X", "label"],
              ["probs", "avgloss"],
              only_loss=only_loss,
              label_prob=label_prob,
              axis=2,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X, label],
              reference=label_softmax_crossent,
          )
  
          self.assertGradientChecks(
              gc, op, [X, label], 0, [1], stepsize=1e-4, threshold=1e-2)
  
      @unittest.skipIf(not workspace.has_gpu_support, "No gpu support")
      @given(**hu.gcs_gpu_only)
      def test_softmax_with_loss_large(self, gc, dc):
          np.random.seed(2603)
          for n in [32]:
              for D in [1000, 2000, 20000]:
                  # n = number of examples, D = |labels|
                  # Initialize X and add 1e-2 for numerical stability
                  X = np.random.rand(n, D).astype(np.float32)
                  X = X + 1e-2
  
                  # Initialize label
                  label = (np.random.rand(n) * D).astype(np.int32)
  
                  # Reference implementation of cross entropy with soft labels
                  def label_softmax_crossent(X, label):
                      probs = np.zeros((n, D))
                      rowmax = np.zeros(n)
                      for i in range(n):
                          rowmax[i] = max(X[i, ])
                          # We need to subtract the max to avoid numerical issues
                          probs[i] = X[i] - rowmax[i]
                          exps = np.exp(probs[i, ])
                          norm = sum(exps)
                          probs[i, ] = exps / norm
  
                      label_xent = [-np.log(max(probs[i][label[i]], 1e-20))
                                    for i in range(n)]
                      avgloss = np.sum(label_xent) / float(n)
                      return (probs, avgloss)
  
                  op = core.CreateOperator(
                      "SoftmaxWithLoss",
                      ["X", "label"],
                      ["probs", "avgloss"]
                  )
  
                  self.assertReferenceChecks(
                      device_option=gc,
                      op=op,
                      inputs=[X, label],
                      reference=label_softmax_crossent,
                  )
  
      @given(n=st.integers(2, 10), D=st.integers(4, 16), **hu.gcs)
      @settings(deadline=None)
      def test_softmax_with_loss_label_prob(self, n, D, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          np.random.seed(2603)
          X = np.random.rand(n, D).astype(np.float32)
          X = X + 1e-2
  
          # Initialize label
          label = np.random.rand(D, n).astype(np.float32)
  
          # normalize labels to sum to 1
          label /= np.sum(label, axis=0)
          label = label.transpose()
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax_crossent(X, label):
              probs = np.zeros((n, D))
              rowmax = np.zeros(n)
              for i in range(n):
                  rowmax[i] = max(X[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
  
              label_xent = np.zeros(X.shape)
              for i in range(n):
                  for j in range(D):
                      label_xent[i][j] = -np.log(
                          max(probs[i, j], 1e-20)) * label[i, j]
              avgloss = np.sum(label_xent) / float(n)
              return (probs, avgloss)
  
          op = core.CreateOperator(
              "SoftmaxWithLoss",
              ["X", "label"],
              ["probs", "avgloss"],
              label_prob=1
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X, label],
              reference=label_softmax_crossent,
          )
  
          self.assertGradientChecks(
              gc, op, [X, label], 0, [1], stepsize=1e-4, threshold=1e-2)
  
      @given(
          n=st.integers(2, 10),
          D=st.integers(4, 16),
          only_loss=st.booleans(),
          **hu.gcs)
      @settings(deadline=None)
      def test_softmax_with_loss_weighted(self, n, D, only_loss, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          np.random.seed(2603)
          X = np.random.rand(n, D).astype(np.float32)
          X = X + 1e-2
  
          # Initialize label
          label = (np.random.rand(n) * D).astype(np.int32)
  
          # Init weights (weight by sample)
          weights = np.random.rand(n).astype(np.float32)
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax_crossent_weighted(X, label, weights):
              probs = np.zeros((n, D))
              rowmax = np.zeros(n)
              for i in range(n):
                  rowmax[i] = max(X[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
  
              label_xent = [-weights[i] * np.log(max(probs[i][label[i]], 1e-20))
                            for i in range(n)]
              avgloss = np.sum(label_xent) / sum(weights)
              return (probs, avgloss)
  
          op = core.CreateOperator(
              "SoftmaxWithLoss",
              ["X", "label", "weights"],
              ["probs", "avgloss"],
              only_loss=only_loss,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X, label, weights],
              reference=label_softmax_crossent_weighted,
          )
  
          self.assertGradientChecks(
              gc, op, [X, label, weights], 0, [1], stepsize=1e-4, threshold=1e-2)
  
      @given(n=st.integers(2, 10), D=st.integers(4, 16), **hu.gcs)
      @settings(deadline=None)
      def test_softmax_with_loss_label_prob_weighted(self, n, D, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          X = np.random.rand(n, D).astype(np.float32)
          X = X + 1e-2
  
          # Initialize label
          label = np.random.rand(D, n).astype(np.float32)
  
          # normalize labels to sum to 1
          label /= np.sum(label, axis=0)
          label = label.transpose()
  
          # Init weights (weight by sample)
          weights = np.random.rand(n).astype(np.float32)
  
          # Reference implementation of cross entropy with soft labels
          def label_softmax_crossent_weighted(X, label, weights):
              probs = np.zeros((n, D))
              rowmax = np.zeros(n)
              for i in range(n):
                  rowmax[i] = max(X[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
  
              label_xent = np.zeros(X.shape)
              for i in range(n):
                  for j in range(D):
                      label_xent[i][j] = -np.log(
                          max(probs[i, j], 1e-20)) * label[i, j] * weights[i]
              avgloss = np.sum(label_xent) / sum(weights)
              return (probs, avgloss)
  
          op = core.CreateOperator(
              "SoftmaxWithLoss",
              ["X", "label", "weights"],
              ["probs", "avgloss"],
              label_prob=1,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[X, label, weights],
              reference=label_softmax_crossent_weighted,
          )
  
          self.assertGradientChecks(
              gc, op, [X, label, weights], 0, [1], stepsize=1e-4, threshold=1e-2)
  
      @given(n=st.integers(2, 5), D=st.integers(2, 4),
             weighted=st.booleans(), **hu.gcs)
      @settings(deadline=None, max_examples=50)
      def test_spatial_softmax_with_loss(self, n, D, weighted, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          W = 18
          H = 12
          np.random.seed(2603)
          X = np.random.rand(n, D, H, W).astype(np.float32)
          X = X + 1e-2
  
          weighted = True
          weights = None
          if weighted:
              weights = np.random.rand(n, H, W).astype(np.float32)
  
          # Initialize label. Some of the labels are (-1), i.e "DONT CARE"
          label = (np.random.rand(n, H, W) * (D + 1)).astype(np.int32) - 1
  
          def label_softmax_crossent_spatial(X, label, weights=None):
              probs = np.zeros((n, D, H, W))
              rowmax = np.zeros((n, H, W))
              label_xent = np.zeros((n, H, W))
              for i in range(n):
                  for x in range(W):
                      for y in range(H):
                          rowmax[i, y, x] = max(X[i, :, y, x])
                          # We need to subtract the max to avoid numerical issues
                          probs[i, :, y, x] = X[i, :, y, x] - rowmax[i, y, x]
                          exps = np.exp(probs[i, :, y, x])
                          probs[i, :, y, x] = exps / sum(exps)
  
                          label_xent[:, y, x] = \
                              [-np.log(max(probs[j, label[i, y, x], y, x], 1e-20))
                               for j in range(n)]
  
              total_xent = 0.0
              total_weight = 0.0
              for y in range(H):
                  for x in range(W):
                      for i in range(n):
                          l = label[i, y, x]
                          if (l != (-1)):
                              w = 1.0 if weights is None else weights[i, y, x]
                              total_xent += \
                                  -np.log(max(probs[i, l, y, x], 1e-20)) * w
                              total_weight += w
              print("Total weight {}".format(total_weight))
  
              return (probs, total_xent / total_weight)
  
          op = core.CreateOperator(
              "SpatialSoftmaxWithLoss",
              ["X", "label"] + ([] if weights is None else ["weights"]),
              ["probs", "avgloss"],
          )
  
          inputs = [X, label] + ([] if weights is None else [weights])
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=inputs,
              reference=label_softmax_crossent_spatial,
          )
  
          self.assertGradientChecks(
              gc, op, inputs, 0, [1], stepsize=1e-4, threshold=1e-2)
  
      @given(n=st.integers(4, 5), D=st.integers(3, 4),
             weighted=st.booleans(), **hu.gcs)
      def test_spatial_softmax_with_loss_allignore(self, n, D, weighted, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          W = 18
          H = 12
          np.random.seed(2603)
          X = np.random.rand(n, D, H, W).astype(np.float32)
          X = X + 1e-2
  
          weighted = True
          weights = None
          if weighted:
              weights = np.random.rand(n, H, W).astype(np.float32)
  
          # Initialize label. All labels as "DONT CARE"
          label = np.zeros((n, H, W)).astype(np.int32) - 1
          print(label)
  
          def label_softmax_crossent_spatial(X, label, weights=None):
              probs = np.zeros((n, D, H, W))
              rowmax = np.zeros((n, H, W))
              label_xent = np.zeros((n, H, W))
              for i in range(n):
                  for x in range(W):
                      for y in range(H):
                          rowmax[i, y, x] = max(X[i, :, y, x])
                          # We need to subtract the max to avoid numerical issues
                          probs[i, :, y, x] = X[i, :, y, x] - rowmax[i, y, x]
                          exps = np.exp(probs[i, :, y, x])
                          probs[i, :, y, x] = exps / sum(exps)
  
                          label_xent[:, y, x] = \
                              [-np.log(max(probs[j, label[i, y, x], y, x], 1e-20))
                              for j in range(n)]
  
              return (probs, 0.0)
  
          op = core.CreateOperator(
              "SpatialSoftmaxWithLoss",
              ["X", "label"] + ([] if weights is None else ["weights"]),
              ["probs", "avgloss"],
          )
  
          inputs = [X, label] + ([] if weights is None else [weights])
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=inputs,
              reference=label_softmax_crossent_spatial,
          )
  
      @given(n=st.integers(4, 5), D=st.integers(3, 4),
             weighted=st.booleans(), **hu.gcs)
      def test_softmax_with_loss_zero_weight(self, n, D, weighted, gc, dc):
          # n = number of examples, D = |labels|
          # Initialize X and add 1e-2 for numerical stability
          np.random.seed(2603)
          X = np.random.rand(n, D).astype(np.float32)
          X = X + 1e-2
  
          weights = np.zeros(n).astype(np.float32)
  
          # Initialize label
          label = (np.random.rand(n) * D).astype(np.int32)
  
          def label_softmax_crossent(X, label, weights=None):
              probs = np.zeros((n, D))
              rowmax = np.zeros((n))
              for i in range(n):
                  rowmax[i] = max(X[i, ])
                  # We need to subtract the max to avoid numerical issues
                  probs[i] = X[i] - rowmax[i]
                  exps = np.exp(probs[i, ])
                  norm = sum(exps)
                  probs[i, ] = exps / norm
              return (probs, 0.0)
  
          op = core.CreateOperator(
              "SoftmaxWithLoss",
              ["X", "label", "weights"],
              ["probs", "avgloss"]
          )
  
          inputs = [X, label] + ([] if weights is None else [weights])
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=inputs,
              reference=label_softmax_crossent,
          )
  
      @unittest.skipIf(not workspace.has_gpu_support, "No gpu support")
      def test_compare_cpugpu(self):
          '''
          Additional test that checks CPU and GPU returns same values
          with larger examples. This is mainly to test the more complex
          GPU implementation is correct.
          '''
          from caffe2.proto import caffe2_pb2
  
          for _j in range(3):
              gpuop = core.CreateOperator(
                  "SpatialSoftmaxWithLoss",
                  ["X_gpu", "label_gpu"],
                  ["probs_gpu", "avgloss_gpu"],
                  device_option=core.DeviceOption(workspace.GpuDeviceType, 0)
              )
  
              cpuop = core.CreateOperator(
                  "SpatialSoftmaxWithLoss",
                  ["X_cpu", "label_cpu"],
                  ["probs_cpu", "avgloss_cpu"],
                  device_option=core.DeviceOption(caffe2_pb2.CPU)
              )
  
              n = 8
              D = 4
              W = 64 + int(np.random.rand(1) * 1024)
              H = 64 + int(np.random.rand(1) * 1024)
  
              print("W: {} H: {}".format(W, H))
  
              X = np.random.rand(n, D, H, W).astype(np.float32)
              X = X + 1e-2
  
              # Initialize label. Some of the labels are (-1), i.e "DONT CARE"
              label = (np.random.rand(n, H, W) * (D + 1)).astype(np.int32) - 1
  
              gpu0 = core.DeviceOption(workspace.GpuDeviceType, 0)
              workspace.FeedBlob("X_cpu", X)
              workspace.FeedBlob("label_cpu", label)
              workspace.FeedBlob("X_gpu", X, device_option=gpu0)
              workspace.FeedBlob("label_gpu", label, device_option=gpu0)
  
              workspace.RunOperatorOnce(gpuop)
              workspace.RunOperatorOnce(cpuop)
  
              probs_gpu = workspace.FetchBlob("probs_gpu")
              probs_cpu = workspace.FetchBlob("probs_cpu")
              loss_gpu = workspace.FetchBlob("avgloss_gpu")
              loss_cpu = workspace.FetchBlob("avgloss_cpu")
  
              np.testing.assert_allclose(probs_gpu, probs_cpu, rtol=1e-4)
              np.testing.assert_allclose(loss_gpu, loss_cpu, rtol=1e-1)
  
  if __name__ == "__main__":
      import unittest
      import random
      random.seed(2603)
      unittest.main()
  