|
- import unittest
-
- import paddle
-
- from fastNLP.core.utils.paddle_utils import paddle_to, paddle_move_data_to_device
-
-
- ############################################################################
- #
- # 测试仅将单个paddle张量迁移到指定设备
- #
- ############################################################################
-
- class PaddleToDeviceTestCase(unittest.TestCase):
- def test_case(self):
- tensor = paddle.rand((4, 5))
-
- res = paddle_to(tensor, "gpu")
- self.assertTrue(res.place.is_gpu_place())
- self.assertEqual(res.place.gpu_device_id(), 0)
- res = paddle_to(tensor, "cpu")
- self.assertTrue(res.place.is_cpu_place())
- res = paddle_to(tensor, "gpu:2")
- self.assertTrue(res.place.is_gpu_place())
- self.assertEqual(res.place.gpu_device_id(), 2)
- res = paddle_to(tensor, "gpu:1")
- self.assertTrue(res.place.is_gpu_place())
- self.assertEqual(res.place.gpu_device_id(), 1)
-
- ############################################################################
- #
- # 测试将参数中包含的所有paddle张量迁移到指定设备
- #
- ############################################################################
-
- class PaddleMoveDataToDeviceTestCase(unittest.TestCase):
-
- def check_gpu(self, tensor, idx):
- """
- 检查张量是否在指定的设备上的工具函数
- """
-
- self.assertTrue(tensor.place.is_gpu_place())
- self.assertEqual(tensor.place.gpu_device_id(), idx)
-
- def check_cpu(self, tensor):
- """
- 检查张量是否在cpu上的工具函数
- """
-
- self.assertTrue(tensor.place.is_cpu_place())
-
- def test_tensor_transfer(self):
- """
- 测试单个张量的迁移
- """
-
- paddle_tensor = paddle.rand((3, 4, 5)).cpu()
- res = paddle_move_data_to_device(paddle_tensor, device=None, data_device=None)
- self.check_cpu(res)
-
- res = paddle_move_data_to_device(paddle_tensor, device="gpu:0", data_device=None)
- self.check_gpu(res, 0)
-
- res = paddle_move_data_to_device(paddle_tensor, device="gpu:1", data_device=None)
- self.check_gpu(res, 1)
-
- res = paddle_move_data_to_device(paddle_tensor, device="gpu:0", data_device="cpu")
- self.check_gpu(res, 0)
-
- res = paddle_move_data_to_device(paddle_tensor, device=None, data_device="gpu:0")
- self.check_gpu(res, 0)
-
- res = paddle_move_data_to_device(paddle_tensor, device=None, data_device="gpu:1")
- self.check_gpu(res, 1)
-
- def test_list_transfer(self):
- """
- 测试张量列表的迁移
- """
-
- paddle_list = [paddle.rand((6, 4, 2)) for i in range(10)]
- res = paddle_move_data_to_device(paddle_list, device=None, data_device="gpu:1")
- self.assertIsInstance(res, list)
- for r in res:
- self.check_gpu(r, 1)
-
- res = paddle_move_data_to_device(paddle_list, device="cpu", data_device="gpu:1")
- self.assertIsInstance(res, list)
- for r in res:
- self.check_cpu(r)
-
- res = paddle_move_data_to_device(paddle_list, device="gpu:0", data_device=None)
- self.assertIsInstance(res, list)
- for r in res:
- self.check_gpu(r, 0)
-
- res = paddle_move_data_to_device(paddle_list, device="gpu:1", data_device="cpu")
- self.assertIsInstance(res, list)
- for r in res:
- self.check_gpu(r, 1)
-
- def test_tensor_tuple_transfer(self):
- """
- 测试张量元组的迁移
- """
-
- paddle_list = [paddle.rand((6, 4, 2)) for i in range(10)]
- paddle_tuple = tuple(paddle_list)
- res = paddle_move_data_to_device(paddle_tuple, device=None, data_device="gpu:1")
- self.assertIsInstance(res, tuple)
- for r in res:
- self.check_gpu(r, 1)
-
- res = paddle_move_data_to_device(paddle_tuple, device="cpu", data_device="gpu:1")
- self.assertIsInstance(res, tuple)
- for r in res:
- self.check_cpu(r)
-
- res = paddle_move_data_to_device(paddle_tuple, device="gpu:0", data_device=None)
- self.assertIsInstance(res, tuple)
- for r in res:
- self.check_gpu(r, 0)
-
- res = paddle_move_data_to_device(paddle_tuple, device="gpu:1", data_device="cpu")
- self.assertIsInstance(res, tuple)
- for r in res:
- self.check_gpu(r, 1)
-
- def test_dict_transfer(self):
- """
- 测试字典结构的迁移
- """
-
- paddle_dict = {
- "tensor": paddle.rand((3, 4)),
- "list": [paddle.rand((6, 4, 2)) for i in range(10)],
- "dict":{
- "list": [paddle.rand((6, 4, 2)) for i in range(10)],
- "tensor": paddle.rand((3, 4))
- },
- "int": 2,
- "string": "test string"
- }
-
- res = paddle_move_data_to_device(paddle_dict, device="gpu:0", data_device=None)
- self.assertIsInstance(res, dict)
- self.check_gpu(res["tensor"], 0)
- self.assertIsInstance(res["list"], list)
- for t in res["list"]:
- self.check_gpu(t, 0)
- self.assertIsInstance(res["int"], int)
- self.assertIsInstance(res["string"], str)
- self.assertIsInstance(res["dict"], dict)
- self.assertIsInstance(res["dict"]["list"], list)
- for t in res["dict"]["list"]:
- self.check_gpu(t, 0)
- self.check_gpu(res["dict"]["tensor"], 0)
-
- res = paddle_move_data_to_device(paddle_dict, device="gpu:0", data_device="cpu")
- self.assertIsInstance(res, dict)
- self.check_gpu(res["tensor"], 0)
- self.assertIsInstance(res["list"], list)
- for t in res["list"]:
- self.check_gpu(t, 0)
- self.assertIsInstance(res["int"], int)
- self.assertIsInstance(res["string"], str)
- self.assertIsInstance(res["dict"], dict)
- self.assertIsInstance(res["dict"]["list"], list)
- for t in res["dict"]["list"]:
- self.check_gpu(t, 0)
- self.check_gpu(res["dict"]["tensor"], 0)
-
- res = paddle_move_data_to_device(paddle_dict, device=None, data_device="gpu:1")
- self.assertIsInstance(res, dict)
- self.check_gpu(res["tensor"], 1)
- self.assertIsInstance(res["list"], list)
- for t in res["list"]:
- self.check_gpu(t, 1)
- self.assertIsInstance(res["int"], int)
- self.assertIsInstance(res["string"], str)
- self.assertIsInstance(res["dict"], dict)
- self.assertIsInstance(res["dict"]["list"], list)
- for t in res["dict"]["list"]:
- self.check_gpu(t, 1)
- self.check_gpu(res["dict"]["tensor"], 1)
-
- res = paddle_move_data_to_device(paddle_dict, device="cpu", data_device="gpu:0")
- self.assertIsInstance(res, dict)
- self.check_cpu(res["tensor"])
- self.assertIsInstance(res["list"], list)
- for t in res["list"]:
- self.check_cpu(t)
- self.assertIsInstance(res["int"], int)
- self.assertIsInstance(res["string"], str)
- self.assertIsInstance(res["dict"], dict)
- self.assertIsInstance(res["dict"]["list"], list)
- for t in res["dict"]["list"]:
- self.check_cpu(t)
- self.check_cpu(res["dict"]["tensor"])
|