device_guard

注意：该API仅支持静态图模式

paddle.static.device_guard ( device=None ) [源代码]

一个用于指定OP运行设备的上下文管理器。

参数：

  • device (str|None) – 指定上下文中使用的设备。它可以是’cpu’或者’gpu‘，当它被设置为’cpu’或者’gpu’时，创建在该上下文中的OP将被运行在CPUPlace或者CUDAPlace上。若设置为’gpu’，同时程序运行在单卡模式下，设备的索引将与执行器的设备索引保持一致。默认值：None，在该上下文中的OP将被自动地分配设备。

代码示例

  1. import paddle
  3. paddle.enable_static()
  4. support_gpu = paddle.is_compiled_with_cuda()
  5. place = paddle.CPUPlace()
  6. if support_gpu:
  7.     place = paddle.CUDAPlace(0)
  9. # if GPU is supported, the three OPs below will be automatically assigned to CUDAPlace(0)
  10. data1 = paddle.full(shape=[1, 3, 8, 8], fill_value=0.5, dtype='float32')
  11. data2 = paddle.full(shape=[1, 3, 64], fill_value=0.5, dtype='float32')
  12. shape = paddle.shape(data2)
  14. with paddle.static.device_guard("cpu"):
  15.     # Ops created here will be placed on CPUPlace
  16.     shape = paddle.slice(shape, axes=[0], starts=[0], ends=[4])
  17. with paddle.static.device_guard('gpu'):
  18.     # if GPU is supported, OPs created here will be placed on CUDAPlace(0), otherwise on CPUPlace
  19.     out = paddle.reshape(data1, shape=shape)
  21. exe = paddle.static.Executor(place)
  22. exe.run(paddle.static.default_startup_program())
  23. result = exe.run(fetch_list=[out])