1.使用as_in_context()
ctx是需要载入的gpu/cpu
# Load data and label is the right context
data = data.as_in_context(ctx)
label = label.as_in_context(ctx)
2.使用copyto(ctx)
将数据copy到ctx
data = data.copyto(ctx)
label = data.copyto(ctx)
3.gluon.utils.split_and_load()
该函数适合于多GPU训练,将数据分割并载入不同GPU
data_list = gluon.data.utils.split_and_load(data, devices)
label_list = gluon.data.utils.split_and_load(label, devices)
如果一个batch不能平均分配给所有的GPU,需要参数even_split = False
4.检测是否有GPU
# Use GPU if available
if len(mx.test_utils.list_gpus())!=0:#return range(0, 4)
ctx=mx.gpu() #default gpu(0)
else:
ctx=mx.cpu()