set_gpu_mode

Applicability

Product

Supported

Atlas A3 training series products/Atlas A3 inference series products

x

Atlas A2 training products/Atlas A2 inference products

x

Atlas 200I/500 A2 inference product

Atlas inference series products

Atlas training products

Description

Schedules AMCT weight quantization to the GPU.

Prototype

1
set_gpu_mode()

Parameters

None

Returns

None

Restrictions

  • The GPU environment is available and CUDA 10.0 is supported. This API does not support GPU card selection. You can select a GPU card by using the CUDA environment variable (CUDA_VISIBLE_DEVICES) or by using the set_device() API of PyCaffe.
  • If neither the set_gpu_mode API nor the set_cpu_mode API is called, AMCT weight quantization is scheduled to the CPU.

Example

1
2
import amct_caffe as amct
amct.set_gpu_mode()