一. 指定一個gpu訓練的兩種方法：

1.程式碼中指定

import torch
torch.cuda.set_device(id)

2.終端中指定

CUDA_VISIBLE_DEVICES=1 python 你的程式

其中id就是你的gpu編號

二. 多gpu並行訓練：

torch.nn.DataParallel(module,device_ids=None,output_device=None,dim=0)

該函式實現了在module級別上的資料並行使用，注意batch size要大於GPU的數量。

引數 ：

module：需要多GPU訓練的網路模型

device_ids： GPU的編號（預設全部GPU，或[0,1],[0,1,2]）

output_device：（預設是device_ids[0])

dim：tensors被分散的維度，預設是0

在儲存模型時要注意要加上"module",例如：

network.module.state_dict()

以上這篇pytorch 指定gpu訓練與多gpu並行訓練示例就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。