20210518 nestedtensor import

Reviewed By: datumbox

Differential Revision: D28521166

fbshipit-source-id: 9761a92aa22c3979a175b8934c5daa4b7d2663af

Author: cpuhrsch

benchmarks/jit_apply.py

@@ -0,0 +1,74 @@
+import torch
+import nestedtensor
+import nestedtensor
+import utils
+
+
+def vmap(fn):
+    def decorator(arg):
+        if torch.is_tensor(arg):
+            return fn(arg)
+        else:
+            def asd(x):
+                return fn(x)
+            return arg.jit_apply(torch.jit.script(asd))
+    return decorator
+
+
+@torch.jit.script
+def my_fun(x):
+    x = x + 1
+    y = x.abs()
+    return y
+
+# print(e)
+
+
+def gen_current():
+    n = nestedtensor.as_nested_tensor(
+        [torch.randn(256, 128).to(device='cuda') for _ in range(128)])
+
+    def _algorithm():
+        n1 = n + 1
+        n1.abs()
+
+    return _algorithm
+
+
+def gen_jit():
+
+    n = nestedtensor._C._ListNestedTensor(
+        [torch.randn(256, 128).to(device='cuda') for _ in range(128)])
+
+    def gen_my_fun(scalar, tensor):
+        @torch.jit.ignore
+        def get_scalar():
+            return scalar
+
+        @torch.jit.ignore
+        def get_tensor():
+            return tensor
+
+        @torch.jit.script
+        def my_fun(x, y):
+            x = x + get_scalar()
+            x = x + get_tensor()
+            y = y + x.abs()
+            return y
+        return my_fun
+    my_fun = gen_my_fun(3.0, torch.randn(1).to(device='cuda'))
+
+    def _algorithm_jit():
+        nestedtensor._C.jit_apply_function((n, n), my_fun)
+
+    return _algorithm_jit
+
+
+if __name__ == "__main__":
+    # print(utils.benchmark_fn(alg, use_cprofile=True))
+    # alg = gen_list_nested_tensor_construction()
+    # print(utils.benchmark_fn(alg))
+    alg1 = gen_current()
+    print(utils.benchmark_fn(alg1))
+    alg2 = gen_jit()
+    print(utils.benchmark_fn(alg2))

benchmarks/linear.py

@@ -0,0 +1,43 @@
+import torch
+import nestedtensor
+import utils
+
+import random
+random.seed(1010)
+
+BDIM=10
+
+# Performance tanks hard for lots of small Tensors as expected
+RAND_INTS = [random.randint(100, 300) for _ in range(BDIM)]
+
+OUTDIM=256
+GOALDIM=512
+
+TENSORS0 = [torch.rand(i, OUTDIM).cuda() for i in RAND_INTS]
+
+def gen_t_linear():
+    nt0 = nestedtensor.nested_tensor(TENSORS0, device=torch.device('cuda'), dtype=torch.float)
+    data, _ = nt0.to_tensor_mask()
+    lin = torch.nn.Linear(OUTDIM, GOALDIM).cuda()
+
+    def t():
+        lin(data)
+    return t
+
+
+@torch.inference_mode()
+def gen_nt_linear():
+    nt0 = nestedtensor.nested_tensor(TENSORS0, device=torch.device('cuda'), dtype=torch.float)
+    lin = torch.nn.Linear(OUTDIM, GOALDIM).cuda()
+
+    def nt():
+        lin(nt0)
+        # print("nt0.size()")
+        # print(nt0.size())
+        # import sys; sys.exit(1)
+    return nt
+
+
+if __name__ == "__main__":
+    print(utils.benchmark_fn(gen_t_linear()))
+    print(utils.benchmark_fn(gen_nt_linear()))

build_tools/travis/after_success.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+# This script is meant to be called by the "after_success" step defined in
+# .travis.yml. See http://docs.travis-ci.com/ for more details.
+
+set -e
+
+if [[ "$COVERAGE" == "true" ]]; then
+    # Ignore codecov failures as the codecov server is not
+    # very reliable but we don't want travis to report a failure
+    # in the github UI just because the coverage report failed to
+    # be published.
+    codecov || echo "codecov upload failed"
+fi

build_tools/travis/install.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+# This script is meant to be called by the "install" step defined in
+# .travis.yml. See http://docs.travis-ci.com/ for more details.
+# The behavior of the script is controlled by environment variabled defined
+# in the .travis.yml in the top level folder of the project.
+
+set -e
+
+echo 'List files from cached directories'
+if [ -d $HOME/download ]; then
+    echo 'download:'
+    ls $HOME/download
+fi
+if [ -d $HOME/.cache/pip ]; then
+    echo 'pip:'
+    ls $HOME/.cache/pip
+fi
+
+# Deactivate the travis-provided virtual environment and setup a
+# conda-based environment instead
+deactivate
+
+# Add the miniconda bin directory to $PATH
+export PATH=/home/travis/miniconda3/bin:$PATH
+echo $PATH
+
+# Use the miniconda installer for setup of conda itself
+pushd .
+cd
+mkdir -p download
+cd download
+if [[ ! -f /home/travis/miniconda3/bin/activate ]]
+then
+    if [[ ! -f miniconda.sh ]]
+    then
+        wget http://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh \
+             -O miniconda.sh
+    fi
+    chmod +x miniconda.sh && ./miniconda.sh -b -f
+    conda update --yes conda
+    echo "Creating environment to run tests in."
+    conda create -n testenv --yes python="$PYTHON_VERSION"
+fi
+cd ..
+popd
+
+# Activate the python environment we created.
+source activate testenv
+
+# Install requirements via pip in our conda environment
+pip install -r requirements.txt
+
+# Install the following only if running tests
+if [[ "$SKIP_TESTS" != "true" ]]; then
+    # PyTorch
+    conda install --yes pytorch torchvision -c pytorch
+
+    # Installation
+    python setup.py install
+fi

build_tools/travis/test_script.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# This script is meant to be called by the "script" step defined in
+# .travis.yml. See http://docs.travis-ci.com/ for more details.
+# The behavior of the script is controlled by environment variabled defined
+# in the .travis.yml in the top level folder of the project.
+
+set -e
+
+python --version
+
+run_tests() {
+    if [[ "$RUN_SLOW" == "true" ]]; then
+        TEST_CMD="py.test --runslow -s -v --cov=nestedtensor --durations=20"
+    else
+        TEST_CMD="py.test -v --cov=nestedtensor --durations=20"
+    fi
+    $TEST_CMD
+}
+
+if [[ "$RUN_FLAKE8" == "true" ]]; then
+    flake8
+fi
+
+if [[ "$SKIP_TESTS" != "true" ]]; then
+    run_tests
+fi

nestedtensor/__init__.py

@@ -14,8 +14,3 @@
 from . import _C
 
 from . import nn
-
-# TODO: https://github.com/pytorch/pytorch/issues/34294
-# torch.cat does not call __torch_function__ properly
-from .nested.nested import _new_torch_stack as stack
-from .nested.nested import _new_torch_cat as cat

nestedtensor/csrc/BinaryOps.cpp

@@ -11,6 +11,22 @@ Tensor NestedTensor_add_Tensor(
   Tensor self;
   Tensor other;
   std::tie(self, other) = _expand_other_as(self_, other_);
+  if (is_nested_tensor_impl(self) && !is_nested_tensor_impl(other) &&
+      get_is_contiguous(self)) {
+    int64_t self_dim = get_dim(self);
+    auto self_opt_sizes = get_opt_sizes(self);
+    if (self_opt_sizes[self_dim - 1] && other.dim() == 1 &&
+        (*(self_opt_sizes[self_dim - 1])) == other.size(0)) {
+      Tensor self_buffer = get_buffer(self);
+      Tensor result_buffer =
+          at::add(self_buffer.reshape({-1, other.size(0)}), other)
+              .reshape({-1});
+      return wrap_buffer(
+          std::move(result_buffer),
+          get_efficient_nested_size(self),
+          get_efficient_nested_stride(self));
+    }
+  }
   return map_nested_tensor(
       [&alpha](Tensor s, Tensor o) { return at::add(s, o, alpha); },
       self,
@@ -97,6 +113,50 @@ Tensor& NestedTensor_div_out(
   return out;
 }
 
+Tensor NestedTensor_floor_divide_Tensor(
+    const Tensor& self_,
+    const Tensor& other_) {
+  Tensor self;
+  Tensor other;
+  std::tie(self, other) = _expand_other_as(self_, other_);
+  return map_nested_tensor(
+      [](Tensor s, Tensor o) { return at::floor_divide(s, o); }, self, other);
+}
+
+Tensor& NestedTensor_floor_divide__Tensor(Tensor& self_, const Tensor& other_) {
+  at::Tensor self;
+  at::Tensor other;
+  std::tie(self, other) = _expand_other_as(self_, other_);
+  apply_nested_tensor(
+      [](Tensor& tensor, const Tensor other) {
+        tensor.floor_divide_(other);
+        return tensor;
+      },
+      self,
+      other);
+  return self_;
+}
+
+Tensor& NestedTensor_floor_divide_out(
+    const Tensor& self,
+    const Tensor& other,
+    Tensor& out) {
+  TORCH_CHECK(
+      is_nested_tensor_impl(out),
+      "NT binary out variant requires NT as out argument.");
+  TORCH_CHECK(
+      is_nested_tensor_impl(out, self, other),
+      "binary_out doesn't support non-NT arguments.")
+  apply_nested_tensor(
+      [](Tensor& self, Tensor& other, Tensor& out) {
+        return at::floor_divide_out(self, other, out);
+      },
+      self,
+      other,
+      out);
+  return out;
+}
+
 Tensor NestedTensor_mul_Tensor(const Tensor& self_, const Tensor& other_) {
   Tensor self;
   Tensor other;
@@ -270,13 +330,32 @@ Tensor& NestedTensor_pow__Tensor(Tensor& self_, const Tensor& other_) {
   return self_;
 }
 
+Tensor NestedTensor_pow_Scalar(const Scalar& base, const Tensor& exponent_) {
+  Tensor exponent = exponent_;
+  return map_nested_tensor(
+      [&base](Tensor exponent) { return at::pow(base, exponent); }, exponent);
+}
+
+Tensor NestedTensor_pow_Tensor_Tensor(
+    const Tensor& self_,
+    const Tensor& other_) {
+  Tensor self;
+  Tensor other;
+  std::tie(self, other) = _expand_other_as(self_, other_);
+  return map_nested_tensor(
+      [](Tensor s, Tensor o) { return at::pow(s, o); }, self, other);
+}
+
 TORCH_LIBRARY_IMPL(aten, NestedTensor, m) {
   nt_impl(m, "add.Tensor", NestedTensor_add_Tensor);
   nt_impl(m, "add_.Tensor", NestedTensor_add__Tensor);
   nt_impl(m, "add.out", NestedTensor_add_out);
   nt_impl(m, "div.Tensor", NestedTensor_div_Tensor);
   nt_impl(m, "div_.Tensor", NestedTensor_div__Tensor);
   nt_impl(m, "div.out", NestedTensor_div_out);
+  nt_impl(m, "floor_divide", NestedTensor_floor_divide_Tensor);
+  nt_impl(m, "floor_divide_.Tensor", NestedTensor_floor_divide__Tensor);
+  nt_impl(m, "floor_divide.out", NestedTensor_floor_divide_out);
   nt_impl(m, "mul.Tensor", NestedTensor_mul_Tensor);
   nt_impl(m, "mul_.Tensor", NestedTensor_mul__Tensor);
   nt_impl(m, "mul.out", NestedTensor_mul_out);
@@ -289,6 +368,8 @@ TORCH_LIBRARY_IMPL(aten, NestedTensor, m) {
   nt_impl(m, "atan2", NestedTensor_atan2);
   nt_impl(m, "remainder.Tensor", NestedTensor_remainder_Tensor);
   nt_impl(m, "pow_.Tensor", NestedTensor_pow__Tensor);
+  nt_impl(m, "pow.Scalar", NestedTensor_pow_Scalar);
+  nt_impl(m, "pow.Tensor_Tensor", NestedTensor_pow_Tensor_Tensor);
 }
 
 } // namespace at

nestedtensor/csrc/activation.cpp

@@ -9,6 +9,12 @@ namespace F = torch::nn::functional;
 namespace at {
 
 Tensor NestedTensor_gelu(const Tensor& self) {
+  if (is_nested_tensor_impl(self) && get_is_contiguous(self)) {
+    return wrap_buffer(
+        at::gelu(get_buffer(self)),
+        get_efficient_nested_size(self),
+        get_efficient_nested_stride(self));
+  }
   return map_nested_tensor(
       [](at::Tensor tensor) { return at::gelu(tensor); }, self);
 }

nestedtensor/csrc/autograd_functions.cpp

@@ -9,8 +9,11 @@ namespace F = torch::nn::functional;
 namespace at {
 
 Tensor NestedTensor_dropout(const Tensor& input, double p, bool train) {
-  return map_nested_tensor(
-      [&](const at::Tensor t) { return at::dropout(t, p, train); }, input);
+  if (train) {
+    return map_nested_tensor(
+        [&](const at::Tensor t) { return at::dropout(t, p, train); }, input);
+  }
+  return input;
 }
 
 Tensor NestedTensor_upsample_bilinear2d(
@@ -100,8 +103,7 @@ Tensor NestedTensor_batch_norm(
     AT_ERROR("running_var must be defined in evaluation mode");
   }
   if (weight) {
-    check_dims_match_num_input_features(
-        "weight", n_input, weight->numel());
+    check_dims_match_num_input_features("weight", n_input, weight->numel());
   }
   if (bias) {
     check_dims_match_num_input_features("bias", n_input, bias->numel());