mxnet學習(6):使用GPU進行訓練

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()
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章