Fast Tensor access in python?

yzhu · February 4, 2017, 7:17am

In lua torch, we can access a Tensor using luajit-FFI pointer as fast as in C. Do we have similar thing in pytorch?

smth · February 4, 2017, 2:58pm

no. Python doesn’t have JITting.

However, you can use Cython, cython code is as concise and nice, and also has indexing and stuff.
For now, in your cython code, you can convert the Tensor to .numpy() and use cython to modify the Tensor.
Here’s an example of using cython to modify numpy arrays:

github.com

rbgirshick/py-faster-rcnn/blob/master/lib/nms/cpu_nms.pyx

# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

import numpy as np
cimport numpy as np

cdef inline np.float32_t max(np.float32_t a, np.float32_t b):
    return a if a >= b else b

cdef inline np.float32_t min(np.float32_t a, np.float32_t b):
    return a if a <= b else b

def cpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh):
    cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0]
    cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1]
    cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2]

This file has been truncated. show original

yzhu · February 4, 2017, 4:12pm

Thanks a lot for the reply and example. I noticed the documentation mentioned cython and numba, this is really sweet.

another maybe more noob question: since we are talking about the access speed here, how’s the overhead of .numpy() operation?

smth · February 5, 2017, 1:01am

the overhead of .numpy() is zero. We dont do any memory copies whatsoever, we just recast the Tensor C struct to a numpy array (the underlying memory is shared between Tensor and numpy array). so it’s free.

apaszke · February 5, 2017, 11:05am

And if you feel comfortable with working with torch tensors, you can check out our ffi examples.