Preparing for PyTorch v1.0.0

Author: jpuigcerver

.travis.yml

@@ -1,5 +1,5 @@
-dist: trusty
-sudo: false
+dist: xenial
+sudo: true
 language: python
 
 os:
@@ -14,9 +14,10 @@ addons:
       - g++-4.9
 
 python:
-  - '2.7'
-  - '3.5'
-  - '3.6'
+  - 2.7
+  - 3.5
+  - 3.6
+  - 3.7
 
 env:
   - CC=gcc-4.9 CXX=g++-4.9
@@ -25,7 +26,7 @@ before_install:
   # Upgrade PIP to latest version, in order to support --progres-bar
   - pip install -U pip
   # Manually install torch, and pybind11 since they are required in the setup
-  - pip install torch==0.4.1 --progress-bar off
+  - pip install torch==1.0.0 --progress-bar off
   - pip install pybind11 --progress-bar off
 
 install:

README.md

@@ -66,7 +66,7 @@ loss4_2 = ctc(x, y, xs, ys)
 
 - C++11 compiler (tested with GCC 4.9).
 - Python: 2.7, 3.5, 3.6, 3.7 (tested with versions 2.7, 3.5 and 3.6).
-- [PyTorch](http://pytorch.org/) >= 0.4.1 (tested with version 0.4.1).
+- [PyTorch](http://pytorch.org/) >= 1.0.0 (tested with version 1.0.0).
 - For GPU support: [CUDA Toolkit](https://developer.nvidia.com/cuda-zone).
 
 ## Installation

setup.py

@@ -57,6 +57,8 @@ def get_cuda_compile_archs(nvcc_flags=None):
         "third-party/warp-ctc/src/ctc_entrypoint.cu",
         "third-party/warp-ctc/src/reduce.cu",
     ]
+    extra_compile_args["cxx"].append("-DWITH_CUDA")
+    extra_compile_args["nvcc"].append("-DWITH_CUDA")
     extra_compile_args["nvcc"].extend(get_cuda_compile_archs())
     Extension = CUDAExtension
 else:
@@ -66,7 +68,7 @@ def get_cuda_compile_archs(nvcc_flags=None):
 
 setup(
     name="torch-baidu-ctc",
-    version="0.1.1",
+    version="0.2.0",
     description="PyTorch bindings for Baidu Warp-CTC",
     long_description=io.open("README.md", "r").read(),
     long_description_content_type="text/markdown",
@@ -102,6 +104,6 @@ def get_cuda_compile_archs(nvcc_flags=None):
         "Topic :: Software Development :: Libraries",
         "Topic :: Software Development :: Libraries :: Python Modules",
     ],
-    setup_requires=["pybind11", "torch>=0.4.1"],
-    install_requires=["pybind11", "torch>=0.4.1"],
+    setup_requires=["pybind11", "torch>=1.0.0"],
+    install_requires=["pybind11", "torch>=1.0.0"],
 )

src/binding.cc

@@ -4,25 +4,29 @@
 #include <sstream>
 #include <string>
 
-#include <torch/torch.h>
+#include <torch/extension.h>
 #include <ctc.h>
 
 #include <ATen/Context.h>
 #include <ATen/CPUGeneral.h>
 #include <ATen/Device.h>
 #include <ATen/DeviceGuard.h>
 
+#ifdef WITH_CUDA
+#include <THC/THC.h>
+#endif
 
-#define CHECK_CONTIGUOUS(x)                                       \
+
+#define CHECK_CONTIGUOUS(x)                                     \
   AT_CHECK((x).is_contiguous(), #x " must be contiguous")
 
 #define CHECK_CPU(x)                                            \
-  AT_CHECK((x).device().type() == at::Device::Type::CPU,        \
+  AT_CHECK((x).device().type() == c10::Device::Type::CPU,       \
            #x " must be located in the CPU")
 
 #define CHECK_CPU_OR_CUDA(x)                                    \
-  AT_CHECK(((x).device().type() == at::Device::Type::CPU ||     \
-            (x).device().type() == at::Device::Type::CUDA),     \
+  AT_CHECK(((x).device().type() == c10::Device::Type::CPU ||    \
+            (x).device().type() == c10::Device::Type::CUDA),    \
            #x " must be located in the CPU or a CUDA device")
 
 #define CHECK_FLOAT(x)                                                  \
@@ -87,14 +91,17 @@ std::tuple<at::Tensor, at::Tensor> ctc_loss(
   ctcOptions ctc_opts;
   memset(&ctc_opts, 0, sizeof(ctcOptions));
   ctc_opts.blank_label = blank_label;
-  if (x.device().type() == at::Device::Type::CPU) {
+  if (x.device().type() == c10::Device::Type::CPU) {
     ctc_opts.loc = CTC_CPU;
     ctc_opts.num_threads = std::max<unsigned int>(at::get_num_threads(), 0);
-  } else {
+#ifdef WITH_CUDA
+  } else if (x.device().type() == c10::Device::Type::CUDA) {
     ctc_opts.loc = CTC_GPU;
-    const auto index = x.device().index();
     ctc_opts.stream =
-        at::globalContext().getCurrentCUDAStreamOnDevice(index).stream();
+        THCState_getCurrentStream(at::globalContext().getTHCState());
+#endif
+  } else {
+    AT_ERROR("ctc_loss not implemented for the given device type");
   }
 
   // Allocate workspace memory
@@ -105,11 +112,11 @@ std::tuple<at::Tensor, at::Tensor> ctc_loss(
           &workspace_size));
 
   at::TensorOptions workspace_opts(x.device());
-  workspace_opts.dtype(at::ScalarType::Byte);
+  workspace_opts = workspace_opts.dtype(at::ScalarType::Byte);
   at::Tensor workspace =
       at::zeros({static_cast<int64_t>(workspace_size * 10)}, workspace_opts);
 
-  at::DeviceGuard device_guard(x.device());
+  c10::DeviceGuard device_guard(x.device());
   CHECK_WARP_CTC_CALL(
       compute_ctc_loss(
           x.data<float>(),

wheels/create_wheels_cpu.sh

@@ -52,7 +52,13 @@ done;
 set +x;
 ODIR="/host/tmp/pytorch_baidu_ctc/whl/cpu";
 mkdir -p "$ODIR";
-cp /tmp/src/dist/*.whl "$ODIR/";
+readarray -t wheels < <(find /tmp/src/dist -name "*.whl");
+for whl in "${wheels[@]}"; do
+  whl_name="$(basename "$whl")";
+  whl_name="${whl_name/-linux/-manylinux1}";
+  cp "$whl" "${ODIR}/${whl_name}";
+done;
+
 echo "================================================================";
-printf "=== %-56s ===\n" "Copied wheels to ${ODIR:5}";
+printf "=== %-56s ===\n" "Copied ${#wheels[@]} wheels to ${ODIR:5}";
 echo "================================================================";

wheels/create_wheels_cuda.sh

@@ -10,9 +10,9 @@ SOURCE_DIR=$(cd $SDIR/.. && pwd);
 
 if [ ! -f /.dockerenv ]; then
   DOCKER_IMAGES=(
-    soumith/manylinux-cuda80
-    soumith/manylinux-cuda90
-    soumith/manylinux-cuda92
+#    soumith/manylinux-cuda80
+#    soumith/manylinux-cuda90
+    soumith/manylinux-cuda100
   );
   for image in "${DOCKER_IMAGES[@]}"; do
     docker run --runtime=nvidia --rm --log-driver none \
@@ -61,7 +61,13 @@ wheels/fix_deps.sh \
   "libcudart.so.${CUDA_VERSION}" \
   "/usr/local/cuda-${CUDA_VERSION}/lib64/libcudart.so.${CUDA_VERSION}";
 
-rm -rf /opt/rh /usr/local/cuda*;
+# Remove CUDA, since all dependencies should be included.
+# TODO: pip package of PyTorch 1.0.0 for CUDA 10 is not well built, we
+# need CUDA installed!
+if [ ${CUDA_VERSION} != "10.0" ]; then
+  rm -rf /opt/rh /usr/local/cuda*;
+fi;
+
 for py in cp27-cp27mu cp35-cp35m cp36-cp36m cp37-cp37m; do
   echo "=== Testing wheel for $py with CUDA ${CUDA_VERSION} ===";
   export PYTHON=/opt/python/$py/bin/python;
@@ -75,7 +81,13 @@ done;
 set +x;
 ODIR="/host/tmp/pytorch_baidu_ctc/whl/${CUDA_VERSION_S}";
 mkdir -p "$ODIR";
-cp /tmp/src/dist/*.whl "$ODIR/";
+readarray -t wheels < <(find /tmp/src/dist -name "*.whl");
+for whl in "${wheels[@]}"; do
+  whl_name="$(basename "$whl")";
+  whl_name="${whl_name/-linux/-manylinux1}";
+  cp "$whl" "${ODIR}/${whl_name}";
+done;
+
 echo "================================================================";
-printf "=== %-56s ===\n" "Copied wheels to ${ODIR:5}";
+printf "=== %-56s ===\n" "Copied ${#wheels[@]} wheels to ${ODIR:5}";
 echo "================================================================";

wheels/install_pytorch_cpu.sh

@@ -8,10 +8,10 @@ PYTHON_VERSIONS=(
   cp37-cp37m
 );
 PYTORCH_WHEELS=(
-  http://download.pytorch.org/whl/cpu/torch-0.4.1-cp27-cp27mu-linux_x86_64.whl
-  http://download.pytorch.org/whl/cpu/torch-0.4.1-cp35-cp35m-linux_x86_64.whl
-  http://download.pytorch.org/whl/cpu/torch-0.4.1-cp36-cp36m-linux_x86_64.whl
-  http://download.pytorch.org/whl/cpu/torch-0.4.1.post2-cp37-cp37m-linux_x86_64.whl
+  http://download.pytorch.org/whl/cpu/torch-1.0.0-cp27-cp27mu-linux_x86_64.whl
+  http://download.pytorch.org/whl/cpu/torch-1.0.0-cp35-cp35m-linux_x86_64.whl
+  http://download.pytorch.org/whl/cpu/torch-1.0.0-cp36-cp36m-linux_x86_64.whl
+  http://download.pytorch.org/whl/cpu/torch-1.0.0-cp37-cp37m-linux_x86_64.whl
 );
 
 for i in $(seq ${#PYTHON_VERSIONS[@]}); do

wheels/install_pytorch_cuda.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-set -ex;
+set -e;
 
 PYTORCH_WHL_PREFIX="http://download.pytorch.org/whl/${CUDA_VERSION_S}";
 PYTHON_VERSIONS=(
@@ -9,10 +9,10 @@ PYTHON_VERSIONS=(
   cp37-cp37m
 );
 PYTORCH_WHEELS=(
-  ${PYTORCH_WHL_PREFIX}/torch-0.4.1-cp27-cp27mu-linux_x86_64.whl
-  ${PYTORCH_WHL_PREFIX}/torch-0.4.1-cp35-cp35m-linux_x86_64.whl
-  ${PYTORCH_WHL_PREFIX}/torch-0.4.1-cp36-cp36m-linux_x86_64.whl
-  ${PYTORCH_WHL_PREFIX}/torch-0.4.1.post2-cp37-cp37m-linux_x86_64.whl
+  ${PYTORCH_WHL_PREFIX}/torch-1.0.0-cp27-cp27mu-linux_x86_64.whl
+  ${PYTORCH_WHL_PREFIX}/torch-1.0.0-cp35-cp35m-linux_x86_64.whl
+  ${PYTORCH_WHL_PREFIX}/torch-1.0.0-cp36-cp36m-linux_x86_64.whl
+  ${PYTORCH_WHL_PREFIX}/torch-1.0.0-cp37-cp37m-linux_x86_64.whl
 );
 
 for i in $(seq ${#PYTHON_VERSIONS[@]}); do