pytorch/caffe2/python/pybind_state_int8.cc
Jerry Zhang aebf3b47ae Remove template parameter from Tensor (#9939)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/9939

Pull Request resolved: https://github.com/facebookresearch/weakly-supervised-action-detection/pull/13

Pull Request resolved: https://github.com/pytorch/translate/pull/166

Pull Request resolved: https://github.com/pytorch/pytorch/pull/9125

Closes https://github.com/pytorch/pytorch/pull/9125

Use inheritance for polymorphism, and remove template parameter
This is to change the templating in call sites, the core implementations will change later

Before Caffe2 Tensor class was compile-time fixed to bind to a particular device/context. With this change, we're making it a runtime property (stored inside the tensor), but preserve the same semantics. For example, one has to specify device type in order to create a Tensor - there are no uninitialized tensors. More specifically the changes are:

1. We added an extra argument *DeviceType* to most of the constructors of the tensor, e.g. (Tensor(DeviceType type)),
2. Semantics of constructor Tensor(const Tensor<SrcContext>& src, ContextForCopy* context); is changed, in this constructor, the second context is passed in to enable us to call the templated Copy function, it could be in a different context as source and target previously, now we'll enforce that the context should have same device type as src, if it is provided.
3. To preserve 'get-or-construct' semantics of Blob, we added specialized getter Blob::GetMutableTensor that verifies both that Blob contains a Tensor and that it's of a correct type
4. Specifically, Tensor type is not default-constructible any more (as we don't have unknown device tensors) and thus some of the code handling STL containers needs to change

Note: Some changes are postponed just to keep this diff a bit smaller. Please see `TODO`s.

Reviewed By: ezyang, houseroad

Differential Revision: D9024330

fbshipit-source-id: e0b8295d2dc6ebe2963383ded5af799ad17164ba
2018-07-27 10:56:39 -07:00

63 lines
2.1 KiB
C++

/**
* Copyright (c) 2016-present, Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Note(jiayq): the import_array function is done inside
// caffe2_python.cc. Read
// http://docs.scipy.org/doc/numpy-1.10.1/reference/c-api.array.html#miscellaneous
// for more details.
#define NO_IMPORT_ARRAY
#include "caffe2/python/pybind_state.h"
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include "caffe2/core/tensor_int8.h"
namespace caffe2 {
namespace python {
class Int8TensorFetcher : public BlobFetcherBase {
public:
pybind11::object Fetch(const Blob& blob) override {
const caffe2::int8::Int8TensorCPU& src =
blob.template Get<caffe2::int8::Int8TensorCPU>();
const int numpy_type = CaffeToNumpyType(src.t.meta());
CAFFE_ENFORCE(numpy_type != -1, "Int8Tensor contains unknown type data");
std::vector<npy_intp> npy_dims;
for (const auto dim : src.t.dims()) {
npy_dims.push_back(dim);
}
auto data_array = pybind11::reinterpret_steal<pybind11::object>(
PyArray_SimpleNew(src.t.dims().size(), npy_dims.data(), numpy_type));
void* ptr = static_cast<void*>(
PyArray_DATA(reinterpret_cast<PyArrayObject*>(data_array.ptr())));
CPUContext context;
context.CopyBytesSameDevice(src.t.nbytes(), src.t.raw_data(), ptr);
context.FinishDeviceComputation();
auto result = pybind11::cast<pybind11::object>(
pybind11::make_tuple(data_array, src.scale, src.zero_point));
return result;
}
};
REGISTER_BLOB_FETCHER(
(TypeMeta::Id<caffe2::int8::Int8TensorCPU>()),
caffe2::python::Int8TensorFetcher);
} // namespace python
} // namespace caffe2