提交 68943f59 编写于 作者: D dangqingqing

Add construction function for LoDTensor in pybind.

上级 e75aab3a
...@@ -125,10 +125,23 @@ The tensor and LoD info should be created before creating the LoDTensor, then ...@@ -125,10 +125,23 @@ The tensor and LoD info should be created before creating the LoDTensor, then
call the set_tensor and set_lod functions to set them. call the set_tensor and set_lod functions to set them.
)DOC") )DOC")
.def("__init__",
[](LoDTensor &instance,
const std::vector<std::vector<size_t>> &lod,
Tensor *t) {
#ifdef PADDLE_ONLY_CPU
new (&instance) LoDTensor(lod, t);
#else
paddle::framework::LoD new_lod;
new_lod.reserve(lod.size());
std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
new (&instance) LoDTensor(new_lod, t);
#endif
})
.def("set_tensor", .def("set_tensor",
[](LoDTensor &self, Tensor *tensor) { self.set_tensor(tensor); }) [](LoDTensor &self, Tensor *tensor) { self.set_tensor(tensor); })
.def("set_lod", .def("set_lod",
[](LoDTensor &self, std::vector<std::vector<size_t>> &lod) { [](LoDTensor &self, const std::vector<std::vector<size_t>> &lod) {
#ifdef PADDLE_ONLY_CPU #ifdef PADDLE_ONLY_CPU
self.set_lod(lod); self.set_lod(lod);
#else #else
......
...@@ -44,61 +44,82 @@ class TestTensor(unittest.TestCase): ...@@ -44,61 +44,82 @@ class TestTensor(unittest.TestCase):
self.assertAlmostEqual(2.0, tensor_array_2[19, 11]) self.assertAlmostEqual(2.0, tensor_array_2[19, 11])
def test_int_lod_tensor(self): def test_int_lod_tensor(self):
scope = core.Scope() places = [core.CPUPlace(), core.GPUPlace(0)]
var = scope.new_var("test_tensor") for place in places:
var_lod = scope.new_var("test_lod_tensor") scope = core.Scope()
place = core.CPUPlace() var = scope.new_var("test_tensor")
var_lod = scope.new_var("test_lod_tensor")
tensor = var.get_tensor() tensor = var.get_tensor()
lod_tensor = var_lod.get_lod_tensor() lod_tensor = var_lod.get_lod_tensor()
tensor.set_dims([4, 4, 6]) tensor.set_dims([4, 4, 6])
tensor.alloc_int(place) tensor.alloc_int(place)
array = numpy.array(tensor) array = numpy.array(tensor)
array[0, 0, 0] = 3 array[0, 0, 0] = 3
array[3, 3, 5] = 10 array[3, 3, 5] = 10
tensor.set(array, place) tensor.set(array, place)
lod_tensor.set_tensor(tensor) lod_tensor.set_tensor(tensor)
lod_tensor.set_lod([[0, 2, 4]]) lod_tensor.set_lod([[0, 2, 4]])
lod_v = numpy.array(lod_tensor.tensor()) lod_v = numpy.array(lod_tensor.tensor())
self.assertTrue(numpy.alltrue(array == lod_v)) self.assertTrue(numpy.alltrue(array == lod_v))
lod = lod_tensor.lod() lod = lod_tensor.lod()
self.assertEqual(0, lod[0][0]) self.assertEqual(0, lod[0][0])
self.assertEqual(2, lod[0][1]) self.assertEqual(2, lod[0][1])
self.assertEqual(4, lod[0][2]) self.assertEqual(4, lod[0][2])
def test_float_lod_tensor(self): def test_float_lod_tensor(self):
places = [core.CPUPlace(), core.GPUPlace(0)]
for place in places:
scope = core.Scope()
var = scope.new_var("test_tensor")
var_lod = scope.new_var("test_lod_tensor")
tensor = var.get_tensor()
lod_tensor = var_lod.get_lod_tensor()
tensor.set_dims([5, 2, 3, 4])
tensor.alloc_float(place)
tensor_array = numpy.array(tensor)
self.assertEqual((5, 2, 3, 4), tensor_array.shape)
tensor_array[0, 0, 0, 0] = 1.0
tensor_array[0, 0, 0, 1] = 2.0
tensor.set(tensor_array, place)
lod_tensor.set_tensor(tensor)
lod_v = numpy.array(lod_tensor.tensor())
self.assertAlmostEqual(1.0, lod_v[0, 0, 0, 0])
self.assertAlmostEqual(2.0, lod_v[0, 0, 0, 1])
self.assertEqual(len(lod_tensor.lod()), 0)
lod_py = [[0, 2, 5], [0, 2, 4, 5]]
lod_tensor.set_lod(lod_py)
lod = lod_tensor.lod()
self.assertListEqual(lod_py, lod)
def test_lod_tensor_init(self):
scope = core.Scope() scope = core.Scope()
var = scope.new_var("test_tensor") var = scope.new_var("test_tensor")
var_lod = scope.new_var("test_lod_tensor")
place = core.CPUPlace() place = core.CPUPlace()
tensor = var.get_tensor() tensor = var.get_tensor()
lod_tensor = var_lod.get_lod_tensor()
tensor.set_dims([5, 2, 3, 4]) tensor.set_dims([5, 2, 3, 4])
tensor.alloc_float(place) tensor.alloc_float(place)
tensor_array = numpy.array(tensor) tensor_array = numpy.array(tensor)
self.assertEqual((5, 2, 3, 4), tensor_array.shape)
tensor_array[0, 0, 0, 0] = 1.0 tensor_array[0, 0, 0, 0] = 1.0
tensor_array[0, 0, 0, 1] = 2.0 tensor_array[0, 0, 0, 1] = 2.0
tensor.set(tensor_array, place) tensor.set(tensor_array, place)
lod_py = [[0, 2, 5], [0, 2, 4, 5]]
lod_tensor.set_tensor(tensor) lod_tensor = core.LoDTensor(lod_py, tensor)
lod_v = numpy.array(lod_tensor.tensor()) lod_v = numpy.array(lod_tensor.tensor())
self.assertAlmostEqual(1.0, lod_v[0, 0, 0, 0]) self.assertAlmostEqual(1.0, lod_v[0, 0, 0, 0])
self.assertAlmostEqual(2.0, lod_v[0, 0, 0, 1]) self.assertAlmostEqual(2.0, lod_v[0, 0, 0, 1])
self.assertEqual(len(lod_tensor.lod()), 0) self.assertListEqual(lod_py, lod_tensor.lod())
lod_py = [[0, 2, 5], [0, 2, 4, 5]]
lod_tensor.set_lod(lod_py)
lod = lod_tensor.lod()
self.assertListEqual(lod_py, lod)
if __name__ == '__main__': if __name__ == '__main__':
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册