device_guard

paddle.device. device_guard ( device ) [源代码]

可以切换当前的 device 为输入指定的 device。

备注

该 API 目前仅支持动态图模式。

参数

  • device (PlaceLike) - 指定的 device,可以是形如 "cpu", "gpu:0" 之类的设备描述字符串,也可以是 paddle.CUDAPlace(0), paddle.CPUPlace() 之类的设备实例。

代码示例

>>> import paddle

>>> # Set the global default device to CPU
>>> paddle.set_device("cpu")
>>> # Temporarily switch to GPU:0 using device_guard with string input
>>> with paddle.device.device_guard("gpu:0"):
...     x = paddle.randn([4, 4])       # Create a Tensor on GPU:0
...     x = x.tanh() * 2               # Perform computation on GPU:0
...     print(x.place)                 # Check the device of the Tensor
Place(gpu:0)

>>> # Set the global default device to GPU:0
>>> paddle.set_device("gpu:0")
>>> # Temporarily switch to CPU using device_guard with Place object (CPUPlace)
>>> cpu_place = paddle.CPUPlace()
>>> with paddle.device.device_guard(cpu_place):
...     x = paddle.randn([4, 4])       # Create a Tensor on CPU
...     x = x.tanh() * 2               # Perform computation on CPU
...     print(x.place)
Place(cpu)