Preparing for PyTorch v1.0.0

.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