set_gpu_mode

功能说明

调用该接口之后，昇腾模型压缩工具执行权重量化的时候，会使用GPU进行加速。

约束说明

用户有GPU环境，且支持CUDA10.0。该接口不支持选择GPU卡，用户可以通过CUDA环境变量(CUDA_VISIBLE_DEVICES)来选择GPU卡，或者使用pycaffe的set_device()接口来选择GPU卡。

函数原型

set_gpu_mode()

参数说明

无。

返回值说明

无。

函数输出

无。

调用示例

import amct_caffe as amct
amct.set_gpu_mode()

如果既不调用set_gpu_mode接口也不调用set_cpu_mode接口。昇腾模型压缩工具默认使用CPU进行权重的量化。