Multi label margin loss

aten/src/ATen/native/LossMulti.h

@@ -0,0 +1,72 @@
+#include <ATen/ATen.h>
+#include <ATen/Dispatch.h>
+#include <ATen/AccumulateType.h>
+
+#pragma once
+
+namespace at { namespace native {
+namespace {
+  static void multilabel_margin_loss_shape_check(
+    int64_t& nframe,
+    int64_t& dim,
+    const int64_t& ndims,
+    TensorArg& target_arg,
+    const Tensor& input,
+    const Tensor& target) {
+    bool valid_inputs = (ndims == 2 && input.size(1) != 0) || (ndims == 1 && input.size(0) != 0) || ndims == 0;
+    TORCH_CHECK(
+                valid_inputs,
+                "Expected non-empty vector or matrix with optional 0-dim batch size, but got: ",
+                input.sizes());
+
+    if (ndims <= 1) {
+      nframe = 1;
+      dim = ndims == 0 ? 1 : input.size(0);
+      TORCH_CHECK(
+                  valid_inputs && target.dim() <= 1 && target.numel() == dim,
+                  "inconsistent size ",
+                  target.sizes(),
+                  " for ",
+                  target_arg);
+    } else {
+      nframe = input.size(0);
+      dim = input.size(1);
+      TORCH_CHECK(
+                  valid_inputs && target.dim() == 2 && target.size(0) == nframe &&
+                  target.size(1) == dim,
+                  "inconsistent size ",
+                  target.sizes(),
+                  " for ",
+                  target_arg);
+    }
+  }
+
+  static void multi_margin_loss_shape_check(
+    int64_t& nframe,
+    int64_t& dim,
+    const int64_t& ndims,
+    TensorArg& target_arg,
+    const Tensor& input,
+    const Tensor& target) {
+    bool valid_inputs = (ndims == 2 && input.size(1) != 0) || (ndims == 1 && input.size(0) != 0) || ndims == 0;
+    if (ndims <= 1) {
+      nframe = 1;
+      dim = ndims == 0 ? 1 : input.size(0);
+    } else {
+      nframe = input.size(0);
+      dim = input.size(1);
+    }
+
+    TORCH_CHECK(
+                valid_inputs,
+                "Expected non-empty vector or matrix with optional 0-dim batch size, but got: ",
+                input.sizes());
+    TORCH_CHECK(
+                valid_inputs && target.dim() <= 1 && target.numel() == nframe,
+                "inconsistent target size, got: ",
+                target.sizes());
+  }
+
+
+}  // anonymous namespace
+}} // namespace at::native

aten/src/ATen/native/LossMultiLabelMargin.cpp

@@ -2,6 +2,7 @@
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
 #include <ATen/TensorUtils.h>
+#include <ATen/native/LossMulti.h>
 
 namespace at {
 namespace native {
@@ -39,6 +40,7 @@ inline scalar_t multilabel_margin_loss_forward_inner_sum_cpu(
       }
     }
   }
+
   return sum;
 }
 
@@ -100,34 +102,32 @@ static void multilabel_margin_loss_forward_out_cpu_template(
     Tensor& is_target,
     int64_t reduction) {
   auto target_arg = TensorArg(target, "target", 2);
-
-  const auto ndims = input.dim();
-
-  TORCH_CHECK(
-      input.numel() > 0 && ndims <= 2,
-      "non-empty vector or matrix expected, got size: ",
-      input.sizes());
-
   int64_t nframe, dim;
+  const int64_t ndims = input.dim();  
   if (ndims <= 1) {
     nframe = 1;
     dim = ndims == 0 ? 1 : input.size(0);
-    TORCH_CHECK(
-        target.numel() > 0 && target.dim() <= 1 && target.numel() == dim,
-        "inconsistent size ",
-        target.sizes(),
-        " for ",
-        target_arg);
-  } else {
+  }
+  else {
     nframe = input.size(0);
     dim = input.size(1);
-    TORCH_CHECK(
-        target.numel() > 0 && target.dim() == 2 && target.size(0) == nframe &&
-            target.size(1) == dim,
-        "inconsistent size ",
-        target.sizes(),
-        " for ",
-        target_arg);
+  }
+  multilabel_margin_loss_shape_check(nframe, dim, ndims, target_arg, input, target);
+
+  // special case target.dim() <= 1: produce scalar output for scalar inputs
+  // even if reduction == Reduction::None
+  if (reduction != Reduction::None || target.dim() <= 1) {
+    output.resize_({});
+  } else {
+    output.resize_({nframe});
+  }
+
+  is_target.resize_as_(target);
+  TORCH_CHECK(is_target.is_contiguous(), "is_target must be contiguous");
+  is_target.zero_();
+
+  if (input.numel() == 0) {
+    return;
   }
 
   TORCH_CHECK(
@@ -138,18 +138,6 @@ static void multilabel_margin_loss_forward_out_cpu_template(
   auto input_contiguous = input.contiguous();
   auto target_contiguous = target.contiguous();
 
-  is_target.resize_as_(target);
-  TORCH_CHECK(is_target.is_contiguous(), "is_target must be contiguous");
-  is_target.zero_();
-
-  // special case target.dim() <= 1: produce scalar output for scalar inputs
-  // even if reduction == Reduction::None
-  if (reduction != Reduction::None || target.dim() <= 1) {
-    output.resize_({});
-  } else {
-    output.resize_({nframe});
-  }
-
   AT_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "multilabel_margin_loss_forward_out_frame", [&] {
         multilabel_margin_loss_forward_out_frame<scalar_t>(
@@ -232,39 +220,22 @@ static void multilabel_margin_loss_backward_out_cpu_template(
     const Tensor& target,
     int64_t reduction,
     const Tensor& is_target) {
+  int64_t nframe, dim;
   CheckedFrom c = "multilabel_margin_loss_backward_cpu_template";
   auto target_arg = TensorArg(target, "target", 3);
   auto is_target_arg = TensorArg(is_target, "is_target", 5);
+  const int64_t ndims = input.dim();
 
-  const auto ndims = input.dim();
-
-  TORCH_CHECK(
-      input.numel() > 0 && ndims <= 2,
-      "non-empty vector or matrix expected, got size: ",
-      input.sizes());
-
-  int64_t nframe, dim;
-  if (ndims <= 1) {
-    nframe = 1;
-    dim = ndims == 0 ? 1 : input.size(0);
-    TORCH_CHECK(
-        target.numel() > 0 && target.dim() <= 1 && target.numel() == dim,
-        "inconsistent size ",
-        target.sizes(),
-        " for ",
-        target_arg);
-  } else {
-    nframe = input.size(0);
-    dim = input.size(1);
-    TORCH_CHECK(
-        target.numel() > 0 && target.dim() == 2 && target.size(0) == nframe &&
-            target.size(1) == dim,
-        "inconsistent size ",
-        target.sizes(),
-        " for ",
-        target_arg);
-  }
+  multilabel_margin_loss_shape_check(nframe, dim, ndims, target_arg, input, target);
   checkSameSize(c, target_arg, is_target_arg);
+
+  grad_input.resize_as_(input);
+  if (grad_input.numel() == 0) {
+    return;
+  }
+
+  TORCH_CHECK(grad_input.is_contiguous(), "grad_input must be contiguous");
+  grad_input.zero_();
 
   TORCH_CHECK(
       target.min().item<int64_t>() >= -1, target_arg, " is out of range");
@@ -275,10 +246,6 @@ static void multilabel_margin_loss_backward_out_cpu_template(
   auto target_contiguous = target.contiguous();
   auto is_target_contiguous = is_target.contiguous();
 
-  grad_input.resize_as_(input);
-  TORCH_CHECK(grad_input.is_contiguous(), "grad_input must be contiguous");
-  grad_input.zero_();
-
   AT_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "multilabel_margin_loss_backward_out_frame", [&] {
         multilabel_margin_loss_backward_out_frame<scalar_t>(

aten/src/ATen/native/LossMultiMargin.cpp

@@ -1,6 +1,7 @@
 #include <ATen/ATen.h>
 #include <ATen/Dispatch.h>
 #include <ATen/AccumulateType.h>
+#include <ATen/native/LossMulti.h>
 
 namespace at {
 namespace native {
@@ -93,34 +94,23 @@ void multi_margin_loss_out_cpu_template(
     Scalar margin,
     const Tensor& weight,
     int64_t reduction) {
+  int64_t nframe, dim;
   const auto ndims = input.dim();
-  TORCH_CHECK(
-      input.numel() > 0 && ndims <= 2,
-      "non-empty vector or matrix expected, got size: ",
-      input.sizes());
+  auto target_arg = TensorArg(target, "target", 2);
 
   TORCH_CHECK(p == 1 || p == 2, "only p == 1 and p == 2 supported");
 
-  int64_t nframe, dim;
-  if (ndims <= 1) {
-    nframe = 1;
-    dim = ndims == 0 ? 1 : input.size(0);
-  } else {
-    nframe = input.size(0);
-    dim = input.size(1);
-  }
-
-  TORCH_CHECK(
-      target.numel() > 0 && target.dim() <= 1 && target.numel() == nframe,
-      "inconsistent target size, got: ",
-      target.sizes());
+  multi_margin_loss_shape_check(nframe, dim, ndims, target_arg, input, target);
 
   // produce a scalar output for 1d input
   if (reduction == Reduction::None && target.dim() > 0) {
     output.resize_({nframe});
   } else {
     output.resize_({});
   }
+  if (input.numel() == 0) {
+    return;
+  }
 
   auto input_contiguous = input.contiguous();
   auto target_contiguous = target.contiguous();
@@ -212,31 +202,20 @@ void multi_margin_loss_backward_out_cpu_template(
     Scalar margin,
     const Tensor& weight,
     int64_t reduction) {
+  int64_t nframe, dim;
+  auto target_arg = TensorArg(target, "target", 2);
   const auto ndims = input.dim();
-  TORCH_CHECK(
-      input.numel() > 0 && ndims <= 2,
-      "non-empty vector or matrix expected, got size: ",
-      input.sizes());
-
+
   TORCH_CHECK(p == 1 || p == 2, "only p == 1 and p == 2 supported");
 
-  int64_t nframe, dim;
-  if (ndims <= 1) {
-    nframe = 1;
-    dim = ndims == 0 ? 1 : input.size(0);
-  } else {
-    nframe = input.size(0);
-    dim = input.size(1);
-  }
-
-  TORCH_CHECK(
-      target.numel() > 0 && target.dim() <= 1 && target.numel() == nframe,
-      "inconsistent target size, got: ",
-      target.sizes());
-
+  multi_margin_loss_shape_check(nframe, dim, ndims, target_arg, input, target);
   grad_input.resize_as_(input);
   TORCH_CHECK(grad_input.is_contiguous(), "grad_input must be contiguous");
 
+  if (input.numel() == 0) {
+    return;
+  }
+
   auto input_contiguous = input.contiguous();
   auto target_contiguous = target.contiguous();
   auto weight_contiguous = weight.contiguous();